Squares of a Sorted List
Problem
A list of integers is sorted in non-decreasing order and may contain negative values. Return a new list holding the square of every value, also in non-decreasing order. Squaring and then sorting takes O(n log n) time; the goal is O(n) by using the order the input already has.
Examples
Input: nums = [-5, -2, 1, 3, 4]
Output: [1, 4, 9, 16, 25]
Why: -5 is the smallest value but gives the largest squareInput: nums = [-7, -3, -2]
Output: [4, 9, 49]
Why: all negative, so the squares come out in reverse orderInput: nums = []
Output: []
Why: edge case, nothing to squareHints
0 / 3
The squares lose their order because large negative values turn into large squares. Where in the input can the biggest square come from?
The largest square always comes from one of the two ends of the input. That suggests two pointers moving inward and an output filled from the back.
Keep one pointer at each end and a write position at the last output slot. Compare the absolute values at the two pointers, write the larger square, move that pointer inward and step the write position back until every slot is filled.
Solution
The largest absolute value is always at one end of the sorted input, so the largest square still unwritten is always one of the two end squares. Two pointers start at the ends and the output is filled from its last slot backwards, the same back-to-front idea used to merge two sorted arrays. Each step writes one square and moves one pointer, so the loop runs exactly n times. Time is O(n) and space is O(n) for the output.
def sorted_squares(nums):
out = [0] * len(nums)
lo, hi = 0, len(nums) - 1
for write in range(len(nums) - 1, -1, -1): # fill from the back
if abs(nums[lo]) > abs(nums[hi]): # the bigger end gives the bigger square
out[write] = nums[lo] * nums[lo]
lo += 1
else:
out[write] = nums[hi] * nums[hi]
hi -= 1
return out
print(sorted_squares([-5, -2, 1, 3, 4])) # -> [1, 4, 9, 16, 25]
print(sorted_squares([-7, -3, -2])) # -> [4, 9, 49]
print(sorted_squares([])) # -> []Stuck on the idea rather than the code? Merge Sorted Arrays covers it.