Squares of a Sorted Array

Patrick Leaton
Problem Description

Given an array of integers A sorted in non-decreasing order, return an array of the squares of each number, also in sorted non-decreasing order.

 

Example 1:

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

Example 2:

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

 

Note:

  1. 1 <= A.length <= 10000
  2. -10000 <= A[i] <= 10000
  3. A is sorted in non-decreasing order.

 

Description taken from https://leetcode.com/problems/squares-of-a-sorted-array/.

Problem Solution

#O(N) Time, O(1) Space
class Solution:
    def sortedSquares(self, A: List[int]) -> List[int]:
        result = []
        right = 0
       
        while right < len(A) and A[right] < 0:
            right += 1
       
        left = right - 1
 
        while left >= 0 and right < len(A):
            if A[left]**2 < A[right]**2:
                result.append(A[left]**2)
                left -= 1
            else:
                result.append(A[right]**2)
                right += 1
       
        while left >= 0:
            result.append(A[left]**2)
            left -= 1
        while right < len(A):
            result.append(A[right]**2)
            right += 1
       
        return result

Problem Explanation


The brute force solution to this problem would be to square all of the values and return the array sorted.  That would require O(N(Log N)) time and constant space.

If we are given a sorted array, many times that is a tip that we can use two pointers.

The whole idea behind this problem is that if a large negative value is squared, it may be the biggest element in the array at that point.  To account for this, we'll need to start with one pointer sitting at the smallest negative value and one pointer sitting at the smallest positive value.

We can solve this problem in linear time by using two pointers moving from the middle of the array outward and within each iteration, we will compare for lesser squared values between the two elements.  Whichever element has the lower square value will be placed into the array first.


Let's start by initializing our result array and right pointer.

        result = []
        right = 0

 

Next, we will use a loop to move the right pointer to the first nonnegative element.  The loop will run while the right pointer is still in-bounds of the array and hasn't reached a nonnegative element.

        while right < len(A) and A[right] < 0:
            right += 1

 

Once we have set our right pointer to its correct location, we will set the left pointer one index to the left of the right pointer.

        left = right - 1

 

Next, we will create a loop to traverse the array that will run while both pointers are still in-bounds.

        while left >= 0 and right < len(A):

 

During each iteration, we will compare the squared elements from the left and right indices.

If the left squared element's value is lower than the right's, then we will append the left squared element to the output array and decrement the left pointer.

            if A[left]**2 < A[right]**2:
                result.append(A[left]**2)
                left -= 1

 

Otherwise, we will append the right squared element and increment the right pointer.

            else:
                result.append(A[right]**2)
                right += 1

 

Once we have broken out of the loop, it means one pointer is out of bounds.  If this happens, we may still have elements that need to be appended to the output array. 

To do this, we will append any remaining values between the left pointer and the start of the array.

We will do the same for the right pointer and the end of the array.

        while left >= 0:
            result.append(A[left]**2)
            left -= 1
        while right < len(A):
            result.append(A[right]**2)
            right += 1

 

Once we have appended each squared input value to the output array in ascending order, we will return it.

        return result


 

Additional Notes

A cleaner solution is to build the array backward, taking the maximum from each iteration instead of the minimum like the previous approach.

The previous two-pointer approach may be more versatile, however.

#O(N) Time, O(N) Space
class Solution:
    def sortedSquares(self, A: List[int]) -> List[int]:
        output = [0] * len(A)
        left, right = 0, len(A) - 1
        for i in range(len(A) -1, -1, -1):
            neg_num = A[left]**2
            pos_num = A[right]**2
            if neg_num > pos_num:
                output[i] = neg_num
                left += 1
            else:
                output[i] = pos_num
                right -= 1
        return output