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:

1
2
3
4

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

Example 2:

1
2
3
4

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

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

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;
    }
};
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

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
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

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;
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

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
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

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
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21

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
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

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.