Find K Closest Elements

Sanjeev SharmaSanjeev Sharma
3 min read

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

Sanjeev Sharma

Written by

Sanjeev Sharma

Full Stack Engineer · E-mopro