Sort Transformed Array — Medium Problem & Solution

You are given a sorted array nums and the quadratic f(x) = a·x² + b·x + c. Apply f to every element and return the results in sorted ascending order.

Difficulty: Medium
Topics: Arrays, Math, Sorting, Two Pointers
Asked at: Amazon, Google, Adobe
Time limit: 2 s
Memory limit: 256 MB
Languages: JavaScript, TypeScript, Python, Java, C++, C, C#, Go, Kotlin, Swift, Rust, PHP and Ruby

Problem statement

You are given a sorted array nums and the quadratic f(x) = a·x² + b·x + c.

Apply f to every element and return the results in sorted ascending order.

Example 1

Input: nums = [-4,-2,2,4], a = 1, b = 3, c = 5
Output: [3,9,15,33]
Explanation: The parabola opens upward, so the extremes of the input give the largest values.

Example 2

Input: nums = [-4,-2,2,4], a = -1, b = 3, c = 5
Output: [-23,-5,1,7]
Explanation: Opening downward reverses which ends are largest.

Example 3

Input: nums = [1,2,3], a = 0, b = 2, c = 1
Output: [3,5,7]
Explanation: With a = 0 the function is linear and increasing.

Constraints

1 <= nums.length <= 200
-100 <= nums[i], a, b, c <= 100
nums is sorted in ascending order.

How to solve Sort Transformed Array

A quadratic is monotone away from its vertex, so on a sorted input the extreme values always sit at the two ends. Which end wins depends on the sign of a, and a two-pointer sweep fills the output in the right direction.

Approach

Evaluate f at the current left and right ends.
If a >= 0 the parabola opens upward: the larger of the two goes at the back of the output, and that pointer moves inward.
If a < 0 it opens downward: the smaller of the two goes at the front, and that pointer moves inward.
Continue until the pointers cross.

Why it works

For a > 0, f decreases then increases, so along a sorted input the maximum over any contiguous stretch is at one of its ends. Repeatedly removing that end therefore emits the values in decreasing order, which filling the output backwards turns into ascending order. The a < 0 case is the mirror image, and a = 0 makes f monotone, which both branches handle correctly.

Complexity

Time — O(n)
Space — O(n) for the output

Pitfalls

Assuming a > 0 always puts the maximum at the right end — with a vertex inside the range, the left end can be larger.
Treating a = 0 as a separate case is unnecessary but harmless; the upward branch handles it.
The values reach about 100 · 100² = 10^6, comfortably inside int.

Reference solution

Python

from typing import List

def sortTransformedArray(nums: List[int], a: int, b: int, c: int) -> List[int]:
    def f(x: int) -> int:
        return a * x * x + b * x + c

    n = len(nums)
    out = [0] * n
    lo, hi = 0, n - 1
    idx = n - 1 if a >= 0 else 0
    while lo <= hi:
        left, right = f(nums[lo]), f(nums[hi])
        if a >= 0:
            if left >= right:
                out[idx] = left
                lo += 1
            else:
                out[idx] = right
                hi -= 1
            idx -= 1
        else:
            if left <= right:
                out[idx] = left
                lo += 1
            else:
                out[idx] = right
                hi -= 1
            idx += 1
    return out

JavaScript

var sortTransformedArray = function(nums, a, b, c) {
    var f = function(x) { return a * x * x + b * x + c; };
    var n = nums.length;
    var out = [];
    for (var t = 0; t < n; t++) out.push(0);
    var lo = 0, hi = n - 1;
    var idx = a >= 0 ? n - 1 : 0;
    while (lo <= hi) {
        var left = f(nums[lo]), right = f(nums[hi]);
        if (a >= 0) {
            if (left >= right) { out[idx] = left; lo++; } else { out[idx] = right; hi--; }
            idx--;
        } else {
            if (left <= right) { out[idx] = left; lo++; } else { out[idx] = right; hi--; }
            idx++;
        }
    }
    return out;
};

Also on the editorial tab: C, C#, C++, Go, Java, Kotlin, PHP, Ruby, Rust, Swift, TypeScript.

All 667 arrays problems · the whole catalogue