Find K Closest Elements

LeetCode 658: Find K Closest Elements Solution & Explanation

Problem Overview

TL;DR: The optimal solution uses binary search to find the starting index of the contiguous subarray of size k by comparing elements at the edges of potential windows.

In "Find K Closest Elements," we are provided with a sorted integer array arr, a target integer x, and an integer k. The objective is to identify k integers from the array that are closest in value to x. If two numbers are equidistant from x, the smaller number is preferred. The final output must be sorted in ascending order.

This is a popular interview question that tests the ability to leverage sorted input data to optimize search operations beyond linear time complexity. The LeetCode 658 problem specifically challenges candidates to handle edge cases where the target x lies outside the array bounds or between existing elements.

Brute Force Approach for Find K Closest Elements

A naive approach ignores the fact that the input array is already sorted. One might calculate the absolute difference between x and every element in arr, sort the entire array based on these differences, and then select the top k elements. Finally, these k elements would need to be sorted back into ascending order to meet the output requirements.

Pseudo-code:

function findClosestElements(arr, k, x):
    pairs = []
    for val in arr:
        diff = abs(val - x)
        pairs.add((diff, val))
    
    sort pairs based on diff, then val
    
    result = []
    for i from 0 to k-1:
        result.add(pairs[i].val)
        
    sort result ascending
    return result

Complexity Analysis: This approach requires sorting the entire array of size $N$. Therefore, the time complexity is $O(N \log N)$. The space complexity is $O(N)$ to store the difference pairs.

Why it fails: While correct, this solution is inefficient. The constraint $N \le 10^4$ implies that an $O(N \log N)$ solution will pass, but it is suboptimal compared to solutions that utilize the pre-sorted nature of arr. In an interview context, ignoring the "sorted" property is a significant red flag.

Algorithm Strategy: Binary Search Patterns

1. Initialize Search Space: Set left = 0 and right = arr.length - k. This covers all valid starting indices for a subarray of size k.
2. Binary Search Loop: While left < right:
   • Calculate mid = left + (right - left) / 2.
   • Compare the distance of x from arr[mid] and arr[mid + k].
   • Check condition: Is x - arr[mid] > arr[mid + k] - x?
     • Note: We assume x >= arr[mid] and arr[mid+k] >= x for the logic, but the absolute difference handles general cases. The simpler logic is: does the element at mid+k offer a smaller distance to x than arr[mid]?
3. Update Pointers:
   • If arr[mid+k] is closer to x (or equal distance but we favor the right side due to logic flow, though ties favor left in this problem), we move left = mid + 1.
   • Otherwise, arr[mid] is closer or equal, so the answer lies at mid or to the left. Set right = mid.
4. Termination: When left == right, we have found the starting index of the k closest elements.
5. Output: Return the subarray from arr[left] to arr[left + k - 1].

Execution Flow

Consider arr = [1, 2, 3, 4, 5], k = 4, x = 3. We need a window of size 4. Valid start indices are 0 to 5-4=1.

1. Initialization:

   • left = 0
   • right = 1

2. Iteration 1:

   • mid = 0 + (1 - 0) / 2 = 0.
   • Compare arr[mid] (which is 1) and arr[mid + k] (which is arr[4] = 5).
   • Distances to x=3:
     • |1 - 3| = 2
     • |5 - 3| = 2
   • Is 2 > 2? False.
   • Since distances are equal, and the problem states "if |a-x| == |b-x| and a < b, a is closer", we prefer 1 over 5.
   • The optimal window is not to the right.
   • Set right = mid = 0.

3. Termination:

   • left is now equal to right (both 0).
   • Loop terminates.

4. Result:

   • Return subarray starting at index 0 with length 4: [1, 2, 3, 4].

Proof of Correctness

The algorithm relies on the convexity of the distance function over a sorted array. The distance |val - x| decreases as val approaches x and increases as val moves away.

When comparing arr[mid] and arr[mid+k], we are essentially checking the gradient of the distance function across the window.

• If arr[mid+k] is closer to x than arr[mid], adding arr[mid+k] and removing arr[mid] improves the solution. This means the optimal start index is strictly greater than mid.
• If arr[mid] is closer or equidistant, moving the window right would increase the total distance deviation (or violate the tie-breaking rule). Thus, the optimal start index is mid or to the left.

Since the decision allows us to eliminate half the search space at each step while maintaining the invariant that the optimal index is within [left, right], the algorithm correctly converges to the optimal starting index.

Pattern Reuse Notes

The binary search pattern used here is specifically about finding a boundary in a derived monotonic property (distance to x). This logic is applicable to:

• LeetCode 34: Find First and Last Position of Element in Sorted Array - Uses the same "Find First/Last Occurrence" subpattern to locate the boundaries of a target value.
• LeetCode 35: Search Insert Position - Finds the first index where a value exceeds or equals the target.

By mastering the "Find First Occurrence" subpattern, you can solve a wide class of problems that require identifying the start of a valid range in logarithmic time.

Reference Implementation

C++ Solution for LeetCode 658

1class Solution {
2public:
3    vector<int> findClosestElements(vector<int>& arr, int k, int x) {
4        // Binary search for the starting index of the window
5        int left = 0;
6        int right = arr.size() - k;
7
8        while (left < right) {
9            int mid = left + (right - left) / 2;
10            
11            // Compare elements just outside the current window overlap
12            // arr[mid] vs arr[mid + k]
13            // If x - arr[mid] > arr[mid + k] - x, it means arr[mid + k] 
14            // is closer to x than arr[mid], so we move right.
15            if (x - arr[mid] > arr[mid + k] - x) {
16                left = mid + 1;
17            } else {
18                right = mid;
19            }
20        }
21
22        // Return the contiguous subarray of size k starting at left
23        return vector<int>(arr.begin() + left, arr.begin() + left + k);
24    }
25};

Java Solution for LeetCode 658

1import java.util.ArrayList;
2import java.util.List;
3
4class Solution {
5    public List<Integer> findClosestElements(int[] arr, int k, int x) {
6        int left = 0;
7        int right = arr.length - k;
8
9        while (left < right) {
10            int mid = left + (right - left) / 2;
11
12            // Check if the element at mid+k is closer to x than the element at mid.
13            // Note: x - arr[mid] represents distance of left-side element
14            // arr[mid + k] - x represents distance of right-side element
15            // If right side is strictly closer, we shift window to the right.
16            if (x - arr[mid] > arr[mid + k] - x) {
17                left = mid + 1;
18            } else {
19                right = mid;
20            }
21        }
22
23        // Construct the result list from the identified window
24        List<Integer> result = new ArrayList<>();
25        for (int i = left; i < left + k; i++) {
26            result.add(arr[i]);
27        }
28        return result;
29    }
30}

Python Solution for LeetCode 658

1class Solution:
2    def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
3        left = 0
4        right = len(arr) - k
5        
6        while left < right:
7            mid = left + (right - left) // 2
8            
9            # Compare distance of arr[mid] and arr[mid + k] to x
10            # If arr[mid + k] is closer, the window must move right
11            if x - arr[mid] > arr[mid + k] - x:
12                left = mid + 1
13            else:
14                right = mid
15                
16        # Return the slice corresponding to the best window
17        return arr[left : left + k]