Отсортировать массив, трансформированный квадратичной функцией

Источник: https://leetcode.com/problems/sort-transformed-array/

Дан отсортированный массив целых чисел nums и три клэффициента a, b and c. Применить квадратичную фугкцию f(x) = ax2 + bx + c к каэжлму элементу nums[i]в массивеи вернуть результат в отсортированном порядк.

Пример 1:

Дано: nums = [-4,-2,2,4], a = 1, b = 3, c = 5
Результат: [3,9,15,33]

Пример 2:

Дано: nums = [-4,-2,2,4], a = -1, b = 3, c = 5
Результат: [-23,-5,1,7]

Идея решения

Квадратичную функцию можно переписать: ax^2+bx+c = a(x+p)^2 + q, где p и q можно выразить через a, b, c.

Но в переписанной форме мы легко можем представить, где распаложена на декартовой плоскости парабола: она будет смещена по горизонтальной оси на -p, и по вертикальной оси на q.

Например: x^2-7x+10=(x-3,5)^2-2,5. То  есть самая нижняя точка параболы смещена вправо на 3,5 и вниз на 2,5, как показано на рисунке:

 

 

Если коэффициент a отрицателен, то парабола будет повернута вниз, а экстремальная  точка будет иметь нибольшее значение y.

 

Мы видим, что значения y будут тем выше/ниже, чем дальше значения x отстоят от нижней/верхней точки параболы. Мы можем идти по массиву в двуз направлениях: слева-направо и справа-налево, рассматривая, какой элемент отстоит от экстремума дальше, и заполнять результат массива с трансформированными значениями от самого большого y до самого наименьшего. И таким образом, результат будет заполняться справа-налево, от самых больших значений, до наименьших.

Имплементация

class Solution {
    public int[] sortTransformedArray(int[] nums, int a, int b, int c) {
       if (nums.length == 0 || nums == null)
            return new int[0];
        int n = nums.length;
        int[] res = new int[n];
        if (a == 0) { // нет квадратичной части,
// элементы в результирующим массиве будут на тех же местах,
// что и в исходном массиве
            for (int i = 0; i < n; i++) {
                // inverse order in case of negative b
                int cur = b >= 0 ? nums[i] : nums[n - 1 - i];
                res[i] = b * cur + c;
            }
            return res;
        }
        // ax^2+bx+c = a(x+p)^2 + q
        // поэтому bx = 2axp
        // поэтому p = b/(2a)
        // a(x+b/(2a))^2+q
        // ПАРАБОЛА
        // смещена по горизонтали на -b/(2a)
        // смещена по вертикали на q (но это нас не интересует)
        double pivot = (double) -b / (2 * a); // смещение по горизонтали

        // считаем, насколько далеко num[i] от экстремума pivot
        // используем два пойнтера: lo и hi
        // чем дальше num[i], тем больше y (если a>0)
        int[] distSorted = new int[n];
        int lo = 0, hi = n - 1, end = n - 1;
        while (lo <= hi) {
            double d1 = pivot - nums[lo], d2 = nums[hi] - pivot;
            // массив расстояний заполняем справа-налево
            if (d1 > d2) {
                distSorted[end--] = nums[lo++];
            } else {
                distSorted[end--] = nums[hi--];
            }
        }
        // если коэффициент a отрицателен, то парабола "смотрит вниз"
        // а дистанции должны идти в обратном порядке
        for (int i = 0; i < n; i++) {
            int cur = a > 0 ? distSorted[i] : distSorted[n - 1 - i];
            res[i] = a * cur * cur + b * cur + c;
        }
        return res;
    }
}