DSA Squares of a Sorted Array

The problem

A list of numbers is sorted, smallest first. The task is to return the squares of its values, also sorted smallest first. Every input value contributes exactly one square.

Squaring is the easy half. Sorting the squares is the problem. A negative value loses its sign when squared, so a very negative value produces a very large square.

cases = [
    ([-4, -1, 0, 3, 10], [0, 1, 9, 16, 100]),
    ([-5, -2, 1, 4], [1, 4, 16, 25]),
    ([0, 1, 2], [0, 1, 4]),
    ([-3, -2, -1], [1, 4, 9]),
]
for nums, expected in cases:
    print("nums", nums, "sorted squares:", expected)

Read the first row slowly. The input -4 squares to 16, and the input 10 squares to 100. The smallest square, 0, comes from the middle of the input, not from either end.

Intuition

The first idea is to square every value and hope the order survives. It does not. Squaring flips the negative half of the list into large positive values, and those values land in the wrong places.

nums = [-4, -1, 0, 3, 10]
squares = [x * x for x in nums]
print("squares in input order:", squares)
print("is sorted:", squares == sorted(squares))

The squares come out as 16, 1, 0, 9, 100. The list falls, then rises. The shape is worth noticing: the small squares sit near zero in the middle, and the large squares sit at the two ends.

That shape is the key observation. At any moment, the largest square not yet placed belongs to one of the two ends of what remains. The middle can never win, because values near zero square to small numbers.

Brute force: square everything, then sort

The straightforward method squares every value into a new list. It then sorts that list from scratch with the built-in sort. The built-in sort is quick, but it starts the ordering problem over. It cannot use the order the input already had.

def brute_squares(nums):
    squares = [x * x for x in nums]
    return sorted(squares)

answer = brute_squares([-4, -1, 0, 3, 10])
print("sorted squares:", answer)

The answer is right. The cost is the sorting pass. Sorting from scratch treats the squares as if nothing were known about them, even though the input arrived in sorted order.

The optimal solution: fill the answer from the back

The largest remaining square is always at the left end or the right end of the input, so the answer is filled from the back, largest first, by comparing the two end squares at each step. That one idea replaces the whole sorting pass with a single scan.

Two pointers mark the ends of the part not yet placed. The bigger end square goes into the last free slot of the answer. The pointer at that end steps inward. The smaller squares near zero are placed last, at the front.

def sorted_squares(nums):
    n = len(nums)
    answer = [0] * n
    left = 0
    right = n - 1
    for pos in range(n - 1, -1, -1):
        if abs(nums[left]) > abs(nums[right]):
            answer[pos] = nums[left] * nums[left]
            left += 1
        else:
            answer[pos] = nums[right] * nums[right]
            right -= 1
    return answer

cases = [
    ([-4, -1, 0, 3, 10], [0, 1, 9, 16, 100]),
    ([-5, -2, 1, 4], [1, 4, 16, 25]),
    ([0, 1, 2], [0, 1, 4]),
    ([-3, -2, -1], [1, 4, 9]),
]
for nums, expected in cases:
    print(nums, "->", sorted_squares(nums), "expected", expected)

All four cases from the start of the article pass. Comparing absolute values is a shortcut for comparing squares. The value with the larger absolute value always has the larger square, so the code can compare before it multiplies.

Walkthrough

Follow [-4, -1, 0, 3, 10] through the loop. The answer starts as five empty slots. Each row below is one placement, from the back slot to the front.

nums = [-4, -1, 0, 3, 10]
n = len(nums)
answer = [0] * n
left = 0
right = n - 1
for pos in range(n - 1, -1, -1):
    left_sq = nums[left] * nums[left]
    right_sq = nums[right] * nums[right]
    if abs(nums[left]) > abs(nums[right]):
        answer[pos] = left_sq
        print("pos", pos, ": left square", left_sq, "beats right square", right_sq, "-> answer", answer)
        left += 1
    else:
        answer[pos] = right_sq
        print("pos", pos, ": right square", right_sq, "beats left square", left_sq, "-> answer", answer)
        right -= 1
print("final:", answer)

The first placement is 100, from the right end. The second is 16, from the left end. The remaining placements then work inward from both sides: 9 from the right end, 1 from the left end, and 0 last of all.

The last placement is 0. Zero sits in the middle of the input and squares to the smallest value. It reaches the answer only when both ends have closed in on it.

Complexity

The two-ends method takes O(n) time: each value is placed exactly once, with one end comparison per placement. It takes O(n) space for the answer list, and O(1) extra space beyond it.

def pair_comparisons(n):
    comparisons = 0
    for i in range(n):
        for j in range(i + 1, n):
            comparisons += 1
    return comparisons

def placement_comparisons(n):
    nums = list(range(-n // 2, n - n // 2))
    left = 0
    right = n - 1
    comparisons = 0
    for pos in range(n - 1, -1, -1):
        comparisons += 1
        if abs(nums[left]) > abs(nums[right]):
            left += 1
        else:
            right -= 1
    return comparisons

for n in [5, 10, 20]:
    print("n =", n, "pair comparisons:", pair_comparisons(n),
          "placement comparisons:", placement_comparisons(n))

A sort that starts from scratch may have to compare any pair of values, and the pair count grows with the square of n. Doubling n roughly quadruples it. The two-ends method makes exactly one comparison per value, so doubling n exactly doubles its work.

Edge cases and pitfalls

An all-positive input is already in square order. The right end wins every placement, and the left pointer never moves. An all-negative input is the mirror image: the left end wins every time, because the most negative value sits there.

Equal end squares are a quiet trap for some versions of this code. When the two ends square to the same value, either end can be placed. The code above places the right end on a tie, because the strict greater-than test falls into the else branch. Both choices give the same sorted answer.

Filling the answer from the front is the classic mistake. The smallest square is not at a known end, so a front-filling version has nothing reliable to place first. Largest first, from the back, is what makes each step a simple choice between two ends.

def sorted_squares(nums):
    n = len(nums)
    answer = [0] * n
    left = 0
    right = n - 1
    for pos in range(n - 1, -1, -1):
        if abs(nums[left]) > abs(nums[right]):
            answer[pos] = nums[left] * nums[left]
            left += 1
        else:
            answer[pos] = nums[right] * nums[right]
            right -= 1
    return answer

for nums in [[], [5], [-7], [1, 2, 3], [-3, -2, -1], [-2, 0, 2], [0, 0, 0]]:
    print(nums, "->", sorted_squares(nums))

The empty list gives an empty answer, because the loop never runs. A single value squares alone. A value and its negative, like -2 and 2, give equal squares side by side.