Description

Given a sorted array of integers nums and integer values a, b and c. Apply a quadratic function of the form f(x) = ax2 + bx + c to each element x in the array.

The returned array must be in sorted order.

Expected time complexity: O(n)

Example

Input: [-4, -2, 2, 4], a = 1, b = 3, c = 5, Output: [3, 9, 15, 33]
Input: [-4, -2, 2, 4], a = -1, b = 3, c = 5, Output: [-23, -5, 1, 7]

Idea

First assume a > 0. Since the array is sorted, and quadratic function first decreases, then increases, the transformed array is unimodal (first drop, then rise).
We find the pivot point, and use two pointers to traverse the left part, which is decreasing, and the right part, which is increasing. Merge them as in the merge sort.
If a < 0, first call the function on -a, -b, -c, then negate and reverse the result.

Solution

class Solution {
public:
    vector<int> sortTransformedArray(vector<int>& nums, int a, int b, int c) {
        int n = nums.size();
        if (a < 0) {
            vector<int> rtn = sortTransformedArray(nums, -a, -b, -c);
            for (auto it = rtn.begin(); it != rtn.end(); it++) *it = -(*it);
            reverse(rtn.begin(), rtn.end());
            return rtn;
        }
        for (int i = 0; i < n; i++) {
            nums[i] = a * nums[i] * nums[i] + b * nums[i] + c;
        }
        int right = 0;
        while (right < n && (right == 0 || nums[right] <= nums[right - 1])) {
            right++;
        }
        int left = right - 1;
        vector<int> rtn;
        while (left >= 0 && right < n) {
            if (nums[left] < nums[right]) {
                rtn.push_back(nums[left--]);
            } else {
                rtn.push_back(nums[right++]);
            }
        }
        while (left >= 0) {
            rtn.push_back(nums[left--]);
        }
        while (right < n) {
            rtn.push_back(nums[right++]);
        }
        return rtn;
    }
};