Squares of a Sorted Array

Given an integer array nums sorted in non-decreasing order, return an array of the squares of each number, also sorted in non-decreasing order. This is LeetCode 977.

Example 1

Input: nums = [-4, -1, 0, 3, 10] Output: [0, 1, 9, 16, 100] Explanation: the squares are [16, 1, 0, 9, 100], which sort to [0, 1, 9, 16, 100].

Example 2

Input: nums = [-7, -3, 2, 3, 11] Output: [4, 9, 9, 49, 121]

How to Think About It

Step 1: The simple way

Square every number and sort the result: O(n log n). It works, but it ignores the fact that the input is already sorted.

Step 2: Why squaring breaks the order

Negative numbers flip: -4 is smaller than 3, but 16 is bigger than 9. A square depends only on the distance from 0, so the squares are smallest near the middle (where the numbers are closest to 0) and largest at the two ends.

Step 3: The biggest square is always at one of the ends

Put lo on the first element and hi on the last. The largest remaining square is one of nums[lo]² and nums[hi]², so:

compare the two squares;
write the larger one into the last free slot of the result;
move that pointer inward.

The result fills from right to left, largest first, and is sorted when the pointers cross.

Walkthrough

nums = [-4, -1, 0, 3, 10].

lo (value)	hi (value)	Squares compared	Written to slot	Result so far
0 (-4)	4 (10)	16 vs 100	4 ← 100	[_, _, _, _, 100]
0 (-4)	3 (3)	16 vs 9	3 ← 16	[_, _, _, 16, 100]
1 (-1)	3 (3)	1 vs 9	2 ← 9	[_, _, 9, 16, 100]
1 (-1)	2 (0)	1 vs 0	1 ← 1	[_, 1, 9, 16, 100]
2 (0)	2 (0)	0 vs 0	0 ← 0	[0, 1, 9, 16, 100]

Java Solution

The loop runs once per slot of the result, from the last slot to the first.

Java

import java.util.Arrays;

public class SquaresOfSortedArray {

    public static int[] sortedSquares(int[] nums) {
        int n = nums.length;
        int[] result = new int[n];
        int lo = 0;
        int hi = n - 1;
        for (int pos = n - 1; pos >= 0; pos--) {   // fill from the largest square down
            int left = nums[lo] * nums[lo];
            int right = nums[hi] * nums[hi];
            if (left > right) {
                result[pos] = left;
                lo++;
            } else {
                result[pos] = right;
                hi--;
            }
        }
        return result;
    }

    public static void main(String[] args) {
        System.out.println(Arrays.toString(sortedSquares(new int[] {-4, -1, 0, 3, 10})));
        // [0, 1, 9, 16, 100]
        System.out.println(Arrays.toString(sortedSquares(new int[] {-7, -3, 2, 3, 11})));
        // [4, 9, 9, 49, 121]
    }
}

The original implementation, SquaresOfSortedArray.java, is on GitHub.

Complexity

Time: O(n): one pass.
Space: O(n) for the result, O(1) besides it.

Filling an array from the back is the same idea as in merge sort's merge step. Two Pointers Basics lists more problems that start from both ends of a sorted array.