Return the k Smallest Squares of a Sorted Integer Array in Order

You are given an integer array nums sorted in non-decreasing order, which may contain negative numbers, and an integer k. Return the k smallest values among the squares of the elements of nums, in non-decreasing order.

With k = len(nums) this returns the squares of every element in sorted order, which was the original question. The follow-up asked for only the k smallest squares and required a solution that uses binary search on the sorted input rather than squaring and sorting every element. Aim for O(log n + k) time.

Function Signature

def k_smallest_squares(nums: list[int], k: int) -> list[int]:

Rules

• Each element contributes exactly one square. Equal squares from different elements (for example from -3 and 3) each appear in the result.
• The result has exactly k values in non-decreasing order. Because it is a sorted list of values, the answer is unique.

Constraints

• 1 <= n <= 10^5, where n = len(nums)
• -10^4 <= nums[i] <= 10^4
• nums is sorted in non-decreasing order.
• 1 <= k <= n
• Every square is at most 10^8, which fits in a 32-bit signed integer.

Examples

Example 1

Input:  nums = [-6, -2, -1, 3, 5], k = 3
Output: [1, 4, 9]

The squares are 36, 4, 1, 9, 25; the three smallest are 1, 4, 9.

Example 2

Input:  nums = [-3, -3, 0, 2, 3], k = 5
Output: [0, 4, 9, 9, 9]

With k = n, every square is returned in sorted order, including all three copies of 9.

Example 3

Input:  nums = [-8, -5, -1], k = 2
Output: [1, 25]

K Smallest Squares of a Sorted Array

You are given an integer array nums sorted in non-decreasing order, which may contain negative numbers, and an integer k. Return the k smallest values among the squares of the elements of nums, as a list in non-decreasing order.

With k = len(nums) the function returns the squares of every element in sorted order. Aim for O(log n + k) time by using binary search on the sorted input rather than squaring and sorting every element.

Function Signature

def k_smallest_squares(nums, k):

Rules

• Each element contributes exactly one square. Equal squares from different elements (for example from -3 and 3, or from repeated values) each appear in the result.
• The result has exactly k values in non-decreasing order. Because it is a sorted list of values, the answer is unique.

Constraints

• 1 <= n <= 10^5, where n = len(nums)
• -10^4 <= nums[i] <= 10^4
• nums is sorted in non-decreasing order.
• 1 <= k <= n
• Every square is at most 10^8, which fits in a 32-bit signed integer. No value can exceed 2^31 - 1, so int is sufficient in Java and C++.

Examples

Example 1

Input:  nums = [-6, -2, -1, 3, 5], k = 3
Output: [1, 4, 9]

The squares are 36, 4, 1, 9, 25; the three smallest are 1, 4, 9.

Example 2

Input:  nums = [-3, -3, 0, 2, 3], k = 5
Output: [0, 4, 9, 9, 9]

With k = n, every square is returned in sorted order, including all three copies of 9.

Hints

1. Think about how squaring changes the order of the negative part of a sorted array compared with the non-negative part.
2. The target is O(log n + k) time, so when k is small you should not need to touch most of the array.
3. Equal squares from different elements, such as -3 and 3 or repeated values, must each appear in the result.

Approach

Binary search finds the split index s, the first position with nums[s] >= 0 (s = n when every element is negative, s = 0 when none is). Squaring is non-increasing on the negative prefix nums[0..s-1] and non-decreasing on the non-negative suffix nums[s..n-1], so reading the prefix from s-1 leftward and the suffix from s rightward yields two non-decreasing sequences of squares. A two-pointer merge of those sequences takes k steps: at each step it compares |nums[left]| = -nums[left] with nums[right] and appends the square of the smaller one, moving that pointer outward. Invariant: after i steps the result holds the i smallest squares in non-decreasing order, because every unconsumed element on either side has absolute value at least as large as its side's current pointer. When the absolute values tie (for example -3 and 3), both squares are equal, so taking either first gives the same value list, and the other is taken on a later step, so every element still contributes exactly one square. When one side is exhausted, the remaining picks come from the other side; since k <= n, the two sides together always hold enough elements. Edge cases: all-negative input (s = n, only the left pointer moves), all-non-negative input (s = 0, only the right pointer moves), runs of zeros (each zero is its own element and contributes a 0), and k = n (the full sorted squares). Every square is at most 10^8, so 32-bit integers suffice.

• Time complexity: O(log n + k)
• Space complexity: O(k)