Squares of a sorted array

Back to index

From the LeetCode problem with the same name. The question statement is:

Given an integer array `nums` sorted in non-decreasing order, return an array of the squares of each number sorted in non-decreasing order.

We are asked to find an algorithm that runs in linear time. The problem is of course that squaring the entries of the input array will put them out of order. Then using a standard sort will probably run worse than linearly. (If we are guaranteed that the number of digits of each number is below some value then I suppose that radix sort might be considered linear.)

So we need to use the structure already existing on the input to help us. In particular: if all the entries happened to be positive then we wouldn't need to sort after squaring. Similarly if all the entries happened to be negative we would just have to reverse the array after sorting.

In the abstract then we are aiming to merge two sorted arrays. This can be done in linear time as shown in the solution below. Note that if this problem had not been specifically about arrays I would have used two stacks rather than keeping track of the leastNegative and leastPositive positions.

class Solution {
  public int[] sortedSquares(int[] nums) {
    int nNegative = 0;
    for (int i = 0; i < nums.length; i++) {
      if (nums[i] < 0) nNegative++;
      nums[i] = nums[i] * nums[i];
    }
    int[] sortedSquares = new int[nums.length];
    int i = 0;
    int leastNegative = nNegative - 1;
    int leastPositive = nNegative;
    while (i < nums.length) {
      if (leastPositive == nums.length) {
        sortedSquares[i] = nums[leastNegative];
        i++;
        leastNegative--;
      } else if (leastNegative == -1) {
        sortedSquares[i] = nums[leastPositive];
        i++;
        leastPositive++;
      } else if (nums[leastNegative] <= nums[leastPositive]) {
        sortedSquares[i] = nums[leastNegative];
        i++;
        leastNegative--;
      } else if (nums[leastNegative] > nums[leastPositive]) {
        sortedSquares[i] = nums[leastPositive];
        i++;
        leastPositive++;
      }
    }
    return sortedSquares;
  }
}