Squares of a Sorted Array (LeetCode 977)

The problem

You get an integer array nums sorted in non-decreasing order. It may contain negatives. Return a new array of each number squared, also sorted in non-decreasing order.

Input:  nums = [-4, -1, 0, 3, 10]
Output: [0, 1, 9, 16, 100]

The squares are 16, 1, 0, 9, 100; sorted, they become 0, 1, 9, 16, 100.

How to think about it

Start with brute force. Square every element and sort the result. It’s one line and runs in O(n log n). It’s correct, but it throws away the fact that the input is already sorted, and the interviewer will ask for O(n).

Why the order breaks. Squaring keeps the order of the non-negatives but reverses the order of the negatives: -4 is the smallest number but has a bigger square than -1. So the squares go down and then up, a “valley” shape.

The key insight: the largest square is always at one of the two ends. The most negative number sits on the far left and the most positive on the far right, and one of them has the biggest absolute value. So compare the two ends, take the bigger square, and place it at the back of the result. Then move that pointer inward and repeat. You’re building the answer from largest to smallest.

Which pattern is this? Converging two pointers. The tell is “sorted input where the extremes matter”, here the two ends compete for the largest value.

Approach

1. Create res of size n, with l = 0 and r = n - 1.
2. Walk a write index k from n - 1 down to 0.
3. If |nums[l]| > |nums[r]|, put nums[l]² at res[k] and move l right.
4. Otherwise put nums[r]² at res[k] and move r left.
5. Return res.

Solution

def sorted_squares(nums):
    n = len(nums); res = [0]*n; l, r = 0, n-1
    for k in range(n-1, -1, -1):              # fill largest slot first
        if abs(nums[l]) > abs(nums[r]):       # bigger magnitude wins slot k
            res[k] = nums[l]*nums[l]; l += 1
        else:
            res[k] = nums[r]*nums[r]; r -= 1
    return res
# sorted_squares([-4,-1,0,3,10]) -> [0,1,9,16,100]

Dry run

On nums = [-4, -1, 0, 3, 10]:

The answer is [0, 1, 9, 16, 100].

Complexity

• Time: O(n). Each step fills one slot and moves one pointer.
• Space: O(n) for the output array. Apart from that, only a few integers.

Edge cases and pitfalls

• All non-negative ([1, 2, 3]): the right pointer wins every time and the order is preserved.
• All negative ([-3, -2, -1]): the left pointer wins every time and the order reverses.
• Single element: one iteration with l == r.
• Filling from the front doesn’t work. The smallest square is somewhere in the middle, not at an end, so you can’t find it with two end pointers.
• Ties (-3 and 3): either side is fine; the other one gets placed next.

Related problems

• Two Sum II – Input Array Is Sorted (LeetCode 167): converging pointers on a sorted array, deciding which end moves.
• Container With Most Water (LeetCode 11): another “compare the two ends, move one” argument.
• Merge Two Sorted Lists (LeetCode 21): this problem is really a merge of two sorted runs (the negatives reversed and the non-negatives).