Find K Closest Elements — LeetCode 658 Heap and Binary Search

Problem Statement#

Given a sorted integer array arr, two integers k and x, return the k closest integers to x in ascending order. Closeness is measured by absolute difference; ties break by smaller value.

Constraints:

• 1 <= k <= arr.length
• 1 <= arr.length <= 10^4
• arr is sorted in ascending order
• -10^4 <= arr[i], x <= 10^4

Input:  arr = [1, 2, 3, 4, 5], k = 4, x = 3
Output: [1, 2, 3, 4]

Input:  arr = [1, 1, 2, 3, 4, 5], k = 4, x = -1
Output: [1, 1, 2, 3]

Why This Problem Matters#

LeetCode 658 Find K Closest Elements is a top Google and Amazon interview problem because it admits two distinct optimal solutions: a size-K max-heap for general arrays and an O(log n + k) binary search on the sliding window boundary, exploiting the sorted property.

Choosing the right approach signals senior judgment: heap for unsorted streams, binary search for sorted arrays. The problem also tests careful tie-breaking (prefer smaller value) — a frequent source of bugs.

Keywords: "K closest interview", "heap binary search hybrid", "FAANG sorted window", "window expansion problem".

The Core Insight#

Because arr is sorted, the answer is always a contiguous window of length K. Find the leftmost valid window with binary search: pick lo such that x - arr[lo] <= arr[lo + k] - x, with ties going left (smaller value).

The heap approach works on any array (sorted or not): push (-distance, value) into a max-heap of capacity K. The bigger heap wins on flexibility; binary search wins on speed.

Visual Dry Run#

arr = [1, 2, 3, 4, 5], k = 4, x = 3. Search for window start in [0, 1].

Window: arr[0..3] = [1, 2, 3, 4].

Solution (Optimal — Binary Search)#

class Solution:
    def findClosestElements(self, arr, k, x):
        lo, hi = 0, len(arr) - k
        while lo < hi:
            mid = (lo + hi) // 2
            if x - arr[mid] > arr[mid + k] - x:
                lo = mid + 1
            else:
                hi = mid
        return arr[lo:lo + k]

var findClosestElements = function(arr, k, x) {
    let lo = 0, hi = arr.length - k;
    while (lo < hi) {
        const mid = (lo + hi) >> 1;
        if (x - arr[mid] > arr[mid + k] - x) lo = mid + 1;
        else hi = mid;
    }
    return arr.slice(lo, lo + k);
};

Time: O(log(n - k) + k) — binary search to find window start, then slice K. Space: O(k) for the result.

Heap Solution (For Unsorted Inputs)#

import heapq
 
class Solution:
    def findClosestElementsHeap(self, arr, k, x):
        h = []
        for v in arr:
            heapq.heappush(h, (-abs(v - x), -v))
            if len(h) > k:
                heapq.heappop(h)
        return sorted(-v for _, v in h)

Time: O(n log k). Space: O(k).

Common Mistakes#

• Forgetting the tie-break: smaller value wins on equal distance.
• Searching hi = len(arr) - 1 instead of len(arr) - k.
• Comparing arr[mid + k - 1] instead of arr[mid + k] — off-by-one.
• Using >= instead of > in the comparison breaks tie-handling.
• Returning unsorted results from the heap path — must sort before returning.

Interview Tips#

• State both approaches and let the interviewer pick.
• For the binary search, use the slick trick: x - arr[mid] > arr[mid + k] - x means the window should shift right.
• For the heap, push (-distance, -value) so larger distance and larger value get evicted first.
• Mention complexity for both; binary search wins on sorted arrays.

Follow-up Questions#

• What if the array is unsorted? Use the heap approach.
• What if K is very large relative to N? Sort by distance and slice K — O(n log n).
• What if elements are streamed? Maintain a size-K max-heap as in solution 2.
• What if you need K closest strings by edit distance? Heap of pairs works the same way.

Key Takeaways#

• LeetCode 658 has two optimal solutions: O(log n + k) binary search and O(n log k) heap.
• The answer is always a contiguous window of length K when the array is sorted.
• Compare x - arr[mid] vs arr[mid + k] - x to decide which side is closer.
• Tie-break by smaller value — push the window left on equal distance.
• Heap approach generalizes to unsorted and streaming inputs.
• Always remember to sort the heap result before returning.
• Off-by-one on arr[mid + k] vs arr[mid + k - 1] is the most common bug.