658. Find K Closest Elements - Explanation

Description

You are given a sorted integer array arr, two integers k and x, return the k closest integers to x in the array. The result should also be sorted in ascending order.

An integer a is closer to x than an integer b if:

|a - x| < |b - x|, or
|a - x| == |b - x| and a < b

Example 1:

Input: arr = [2,4,5,8], k = 2, x = 6

Output: [4,5]

Example 2:

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

Output: [2,3,4]

Constraints:

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

Company Tags

Please upgrade to NeetCode Pro to view company tags.

1. Sorting (Custom Comparator)

Intuition

The problem asks for the k elements closest to x. A straightforward approach is to sort all elements by their distance to x. If two elements have the same distance, we prefer the smaller one. After sorting, we simply take the first k elements and return them in sorted order.

Algorithm

Sort the array using a custom comparator that compares elements by their absolute difference from x. If two elements have the same distance, the smaller element comes first.
Take the first k elements from the sorted array.
Sort these k elements in ascending order (since the output must be sorted).
Return the result.

class Solution:
    def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
        arr.sort(key=lambda num: (abs(num - x), num))
        return sorted(arr[:k])

Time & Space Complexity

Time complexity: $O(n \log n + k \log k)$
Space complexity:
- $O(1)$ or $O(n)$ space depending on the sorting algorithm.
- $O(k)$ space for the output array.

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

2. Linear Scan + Two Pointers

Intuition

Since we need elements closest to x, we can first find the element nearest to x using a linear scan. Once we have this starting point, we expand outward using two pointers, picking the closer element at each step until we have k elements.

Algorithm

Scan through the array to find the index of the element closest to x.
Initialize the result with that element.
Set two pointers: l pointing to the index before the closest element, and r pointing to the index after.
While we have fewer than k elements:
- If both pointers are valid, compare distances and add the closer element.
- If only the left pointer is valid, add from the left.
- If only the right pointer is valid, add from the right.
Sort the result and return it.

class Solution:
    def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
        n = len(arr)
        idx = 0
        for i in range(1, n):
            if abs(x - arr[idx]) > abs(x - arr[i]):
                idx = i

        res = [arr[idx]]
        l, r = idx - 1, idx + 1

        while len(res) < k:
            if l >= 0 and r < n:
                if abs(x - arr[l]) <= abs(x - arr[r]):
                    res.append(arr[l])
                    l -= 1
                else:
                    res.append(arr[r])
                    r += 1
            elif l >= 0:
                res.append(arr[l])
                l -= 1
            elif r < n:
                res.append(arr[r])
                r += 1

        return sorted(res)

Time & Space Complexity

Time complexity: $O(n + k \log k)$
Space complexity:
- $O(1)$ or $O(k)$ space depending on the sorting algorithm.
- $O(k)$ space for the output array.

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

3. Two Pointers

Intuition

Since the array is already sorted, the k closest elements must form a contiguous subarray. We can use two pointers starting at the ends of the array and shrink the window by removing the element that is farther from x until only k elements remain.

Algorithm

Initialize l = 0 and r = n - 1.
While r - l >= k:
- Compare the distances of arr[l] and arr[r] from x.
- Remove the element that is farther by moving the corresponding pointer inward.
- If distances are equal, prefer the left element (smaller value), so move r inward.
Return the subarray from index l to r (inclusive).

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

        return arr[l: r + 1]

Time & Space Complexity

Time complexity: $O(n - 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.

4. Binary Search + Two Pointers

Intuition

Instead of scanning linearly to find the starting point, we can use binary search to quickly locate where x would fit in the sorted array. From that position, we expand outward with two pointers to collect k closest elements.

Algorithm

Use binary search to find the leftmost position where arr[mid] >= x.
Initialize two pointers: l pointing one position to the left of this index, and r at this index.
While we have fewer than k elements (i.e., r - l - 1 < k):
- If l < 0, expand right.
- If r >= n, expand left.
- Otherwise, compare distances and expand toward the closer element.
Return the subarray from index l + 1 to r - 1 (exclusive bounds).

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

        l, r = l - 1, l
        while r - l - 1 < k:
            if l < 0:
                r += 1
            elif r >= len(arr):
                l -= 1
            elif abs(arr[l] - x) <= abs(arr[r] - x):
                l -= 1
            else:
                r += 1

        return arr[l + 1:r]

Time & Space Complexity

Time complexity: $O(\log n + 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.

5. Binary Search

Intuition

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.

Algorithm

Set l = 0 and r = n - k (the rightmost valid starting index).
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.
Return the subarray starting at index l with length k.

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]

Time & Space 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.

Common Pitfalls

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.