Find K Closest Elements
Advertisement
Problem
Given a sorted array, two integers k and x, return the k closest elements to x. In case of ties, prefer the smaller element.
Key insight: Binary search for the left boundary of the k-element window. Check if x - arr[mid] <= arr[mid+k] - x.
Approach — Binary Search on Window
Find the leftmost index of the k-element window. lo = 0, hi = n - k.
Solutions
// C
int* findClosestElements(int* arr, int n, int k, int x, int* retSize) {
int lo = 0, hi = n - k;
while (lo < hi) {
int mid = (lo + hi) / 2;
if (x - arr[mid] > arr[mid + k] - x) lo = mid + 1;
else hi = mid;
}
*retSize = k;
return arr + lo;
}
// C++
vector<int> findClosestElements(vector<int>& arr, int k, int x) {
int lo = 0, hi = (int)arr.size() - k;
while (lo < hi) {
int mid = (lo + hi) / 2;
if (x - arr[mid] > arr[mid + k] - x) lo = mid + 1;
else hi = mid;
}
return vector<int>(arr.begin() + lo, arr.begin() + lo + k);
}
// Java
public List<Integer> findClosestElements(int[] arr, int k, int x) {
int lo = 0, hi = arr.length - k;
while (lo < hi) {
int mid = (lo + hi) / 2;
if (x - arr[mid] > arr[mid + k] - x) lo = mid + 1;
else hi = mid;
}
List<Integer> res = new ArrayList<>();
for (int i = lo; i < lo + k; i++) res.add(arr[i]);
return res;
}
// JavaScript
function findClosestElements(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);
}
# Python
def findClosestElements(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]
Complexity
- Time: O(log(n-k) + k)
- Space: O(1)
Key Insight
We're searching for the start of a window of size k. If x - arr[mid] > arr[mid+k] - x, the right side of the window is closer — shift window right.
Advertisement
← Previous
Kth Largest Element in an ArrayNext →
Relative Ranks