Find K Closest Elements

How it works

We can binary search directly for the starting index of the k-length window. For any starting index m, we compare the distances of arr[m] and arr[m + k] to x. If arr[m + k] is closer, the window should shift right; otherwise, it should stay or shift left. This narrows down the optimal starting position.

Steps

1. Set l = 0 and r = n - k (the rightmost valid starting index).
2. While l < r:
   • Compute m = (l + r) / 2.
   • Compare x - arr[m] with arr[m + k] - x.
   • If x - arr[m] > arr[m + k] - x, the window is too far left, so set l = m + 1.
   • Otherwise, set r = m.
3. Return the subarray starting at index l with length k.

Complexity

• Time complexity: O(log (n - k) + k)
• Space complexity: O(k) for the output array.

Where n is the size of the input array and k is the number of closest elements to find.

Things to watch for

Forgetting the Tie-Breaking Rule

When two elements have the same distance to x, the problem requires preferring the smaller element. A common mistake is treating equal distances arbitrarily or preferring the larger element. This affects both the custom comparator approach and the two-pointer shrinking logic where you should shrink from the right (larger values) when distances are equal.

Not Returning Results in Sorted Order

The output must be sorted in ascending order. When using approaches that collect elements based on distance (like sorting by distance or expanding from a center point), the collected elements may not be in sorted order. Forgetting to sort the final result before returning leads to incorrect output.

Incorrect Binary Search Bounds

In the optimized binary search approach, the search range should be [0, n - k] for the starting index of the window, not [0, n - 1]. Using incorrect bounds can cause the window to extend past the array end or miss valid starting positions, leading to index out of bounds errors or wrong answers.

from typing import List

class Solution:
    def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
        l, r = 0, len(arr) - k
        while l < r:
            m = (l + r) // 2
            if x - arr[m] > arr[m + k] - x:
                l = m + 1
            else:
                r = m
        return arr[l:l + k]

class Solution {
public:
    vector<int> findClosestElements(vector<int>& arr, int k, int x) {
        int l = 0, r = arr.size() - k;
        while (l < r) {
            int m = (l + r) / 2;
            if (x - arr[m] > arr[m + k] - x) {
                l = m + 1;
            } else {
                r = m;
            }
        }
        return vector<int>(arr.begin() + l, arr.begin() + l + k);
    }
};