Sort Transformed Array Problem

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_2 + bx + c to each element x in the array.

The returned array must be in sorted order.

Expected time complexity: O(n)

Examples

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]

Solution

Method 1 – Two Pointers for Quadratic Transformation

Intuition

The function f(x) = ax^2 + bx + c is convex (a > 0) or concave (a < 0). For sorted input, we can use two pointers to fill the result from the correct end, depending on the sign of a.

Approach

1. Use two pointers (left, right) at the ends of the array.
2. For a >= 0, fill result from the end (largest values at ends), for a < 0, fill from the start.
3. At each step, compare f(nums[left]) and f(nums[right]) and place appropriately.

Code

C++

class Solution {
public:
    vector<int> sortTransformedArray(vector<int>& nums, int a, int b, int c) {
        int n = nums.size();
        vector<int> res(n);
        auto f = [&](int x) { return a*x*x + b*x + c; };
        int l = 0, r = n-1, idx = a >= 0 ? n-1 : 0;
        while (l <= r) {
            int lv = f(nums[l]), rv = f(nums[r]);
            if (a >= 0) {
                if (lv > rv) res[idx--] = lv, l++;
                else res[idx--] = rv, r--;
            } else {
                if (lv < rv) res[idx++] = lv, l++;
                else res[idx++] = rv, r--;
            }
        }
        return res;
    }
};

Go

func sortTransformedArray(nums []int, a, b, c int) []int {
    n := len(nums)
    res := make([]int, n)
    f := func(x int) int { return a*x*x + b*x + c }
    l, r := 0, n-1
    idx := n-1
    if a < 0 { idx = 0 }
    for l <= r {
        lv, rv := f(nums[l]), f(nums[r])
        if a >= 0 {
            if lv > rv { res[idx] = lv; idx--; l++ } else { res[idx] = rv; idx--; r-- }
        } else {
            if lv < rv { res[idx] = lv; idx++; l++ } else { res[idx] = rv; idx++; r-- }
        }
    }
    return res
}

Java

class Solution {
    public int[] sortTransformedArray(int[] nums, int a, int b, int c) {
        int n = nums.length;
        int[] res = new int[n];
        int l = 0, r = n-1, idx = a >= 0 ? n-1 : 0;
        while (l <= r) {
            int lv = a*nums[l]*nums[l] + b*nums[l] + c;
            int rv = a*nums[r]*nums[r] + b*nums[r] + c;
            if (a >= 0) {
                if (lv > rv) res[idx--] = lv; else res[idx--] = rv;
                if (lv > rv) l++; else r--;
            } else {
                if (lv < rv) res[idx++] = lv; else res[idx++] = rv;
                if (lv < rv) l++; else r--;
            }
        }
        return res;
    }
}

Kotlin

class Solution {
    fun sortTransformedArray(nums: IntArray, a: Int, b: Int, c: Int): IntArray {
        val n = nums.size
        val res = IntArray(n)
        var l = 0; var r = n-1
        var idx = if (a >= 0) n-1 else 0
        fun f(x: Int) = a*x*x + b*x + c
        while (l <= r) {
            val lv = f(nums[l]); val rv = f(nums[r])
            if (a >= 0) {
                if (lv > rv) { res[idx--] = lv; l++ } else { res[idx--] = rv; r-- }
            } else {
                if (lv < rv) { res[idx++] = lv; l++ } else { res[idx++] = rv; r-- }
            }
        }
        return res
    }
}

Python

class Solution:
    def sortTransformedArray(self, nums: list[int], a: int, b: int, c: int) -> list[int]:
        n = len(nums)
        res = [0]*n
        l, r = 0, n-1
        idx = n-1 if a >= 0 else 0
        def f(x): return a*x*x + b*x + c
        while l <= r:
            lv, rv = f(nums[l]), f(nums[r])
            if a >= 0:
                if lv > rv:
                    res[idx] = lv; idx -= 1; l += 1
                else:
                    res[idx] = rv; idx -= 1; r -= 1
            else:
                if lv < rv:
                    res[idx] = lv; idx += 1; l += 1
                else:
                    res[idx] = rv; idx += 1; r -= 1
        return res

Rust

impl Solution {
    pub fn sort_transformed_array(nums: Vec<i32>, a: i32, b: i32, c: i32) -> Vec<i32> {
        let n = nums.len();
        let mut res = vec![0; n];
        let (mut l, mut r) = (0, n-1);
        let mut idx = if a >= 0 { n-1 } else { 0 };
        let f = |x: i32| a*x*x + b*x + c;
        while l <= r {
            let lv = f(nums[l]);
            let rv = f(nums[r]);
            if a >= 0 {
                if lv > rv { res[idx] = lv; idx -= 1; l += 1; }
                else { res[idx] = rv; idx -= 1; r -= 1; }
            } else {
                if lv < rv { res[idx] = lv; idx += 1; l += 1; }
                else { res[idx] = rv; idx += 1; r -= 1; }
            }
        }
        res
    }
}

TypeScript

class Solution {
    sortTransformedArray(nums: number[], a: number, b: number, c: number): number[] {
        const n = nums.length;
        const res = Array(n).fill(0);
        let l = 0, r = n-1;
        let idx = a >= 0 ? n-1 : 0;
        const f = (x: number) => a*x*x + b*x + c;
        while (l <= r) {
            const lv = f(nums[l]), rv = f(nums[r]);
            if (a >= 0) {
                if (lv > rv) { res[idx--] = lv; l++; }
                else { res[idx--] = rv; r--; }
            } else {
                if (lv < rv) { res[idx++] = lv; l++; }
                else { res[idx++] = rv; r--; }
            }
        }
        return res;
    }
}

Complexity

• ⏰ Time complexity: O(n) where n is the length of nums.
• 🧺 Space complexity: O(n) for the result array.