Return the k Smallest Squares of a Sorted Integer Array in Order
Company: Uber
Role: Software Engineer
Category: Coding & Algorithms
Difficulty: hard
Interview Round: Onsite
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
```python
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**
```text
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**
```text
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**
```text
Input: nums = [-8, -5, -1], k = 2
Output: [1, 25]
```
Overview: A coding problem on a sorted integer array that may contain negative numbers, asking for the k smallest squares of its elements in sorted order. It tests exploiting sorted input, reasoning about how squaring reorders negative and positive values, and meeting a logarithmic-plus-k time target.
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
```python
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**
```text
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**
```text
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`.
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 exceeds 2^31 - 1.
Examples
Input: ([-6, -2, -1, 3, 5], 3)
Expected Output: [1, 4, 9]
Explanation: Source Example 1: squares 36, 4, 1, 9, 25; the three smallest are 1, 4, 9.
Input: ([-3, -3, 0, 2, 3], 5)
Expected Output: [0, 4, 9, 9, 9]
Explanation: Source Example 2: k = n returns every square sorted, keeping all three copies of 9.
Hints
- Think about how squaring changes the order of the negative part of a sorted array compared with the non-negative part.
- The target is O(log n + k) time, so when k is small you should not need to touch most of the array.
- Equal squares from different elements, such as -3 and 3 or repeated values, must each appear in the result.