360. Sort Transformed Array

• Difficulty: Medium

• Topics: Math, Two Pointers

• Similar Questions:

  • Squares of a Sorted Array

Problem:

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

The returned array must be in sorted order.

Expected time complexity: O(n)

Example 1:

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

Example 2:

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

Solutions:

class Solution {
public:
    vector<int> sortTransformedArray(vector<int>& nums, int a, int b, int c) {
        this->a = a;
        this->b = b;
        this->c = c;

        if (nums.size() == 0)   return {};
        if (a == 0) {
            if (b < 0) {
                reverse(nums.begin(), nums.end());
            }
            vector<int> ret;
            for (auto& val : nums) {
                ret.push_back(f(val));
            }

            return ret;
        }

        double vertex = -b / (double)(2 * a);

        int left = 0, right = nums.size() - 1; // largest value smaller than

        while (left < right) {
            int mid = right - (right - left) / 2;
            if (vertex <= nums[mid]) {
                right = mid - 1;
            } else {
                left = mid;
            }
        }

        vector<int> ret;
        right = left + 1;
        while (left >= 0 && right < nums.size()) {
            if (abs(nums[left] - vertex) <= abs(nums[right] - vertex)) {
                ret.push_back(f(nums[left]));
                --left;
            } else {
                ret.push_back(f(nums[right]));
                ++right;
            }
        }

        if (left < 0) {
            while (right < nums.size()) {
                ret.push_back(f(nums[right++]));
            }
        }

        if (right >= nums.size()) {
            while (left >= 0) {
                ret.push_back(f(nums[left--]));
            }
        }

        if (a < 0) {
            reverse(ret.begin(), ret.end());
        }

        return ret;
    }

private:
    inline int f(int x) {
        return a * x * x + b * x + c;
    }

    int a;
    int b;
    int c;

};