Find K Closest Elements

The problem. Given a sorted array arr, and integers k and x, return the k values closest to x, sorted in ascending order. Value a is closer than b if |a − x| < |b − x|, or if they're equally far and a < b.

Example

arr = [1, 2, 3, 4, 5], k = 4, x = 3      ->  [1, 2, 3, 4]
arr = [1, 1, 2, 3, 4, 5], k = 4, x = -1  ->  [1, 1, 2, 3]
arr = [1, 3, 5, 8, 9, 12], k = 3, x = 7  ->  [5, 8, 9]

The idea: binary search the window's start

Because the array is sorted, the k closest values are always k neighbours — a window arr[s..s+k-1]. So the real question is the start s, somewhere in 0..n−k. Compare window s with window s + 1: they share everything except arr[s] (only in the first) and arr[s + k] (only in the second). If x − arr[s] > arr[s + k] − x, then arr[s] is farther away than arr[s + k], so the window should start later. Otherwise start s is at least as good. Those answers go "later, later, …, here or earlier" — binary search finds the first start that's good enough.

1. Sort by distance

1function findClosestElements(arr, k, x) {2  const byDistance = [...arr].sort((a, b) => Math.abs(a - x) - Math.abs(b - x) || a - b);   // ties: smaller first3  return byDistance.slice(0, k).sort((a, b) => a - b);4}

Time O(n log n), space O(n). It throws away the fact that arr is already sorted.

2. Shrink a window from both ends

1function findClosestElements(arr, k, x) {2  let left = 0;3  let right = arr.length - 1;4  while (right - left + 1 > k) {5    if (x - arr[left] <= arr[right] - x) right--;   // the right end is no closer: drop it (ties keep the smaller)6    else left++;                                   // the left end is farther: drop it7  }8  return arr.slice(left, right + 1);9}

Time O(n − k), space O(1) besides the output. The farther of the two ends can never be among the k closest while the other end is still in.

3. Binary search for the window's start

1function findClosestElements(arr, k, x) {2  let low = 0;3  let high = arr.length - k;                 // the window arr[s..s+k-1] starts somewhere in 0..n-k4  while (low < high) {5    const mid = low + Math.floor((high - low) / 2);6    if (x - arr[mid] > arr[mid + k] - x) low = mid + 1;   // arr[mid + k] is closer: start later7    else high = mid;                                     // arr[mid] is at least as close: start here or earlier8  }9  return arr.slice(low, low + k);10}

Time O(log(n − k) + k), space O(1) besides the output. Compare the signed differences, not absolute values: with duplicates, |arr[mid] − x| > |arr[mid + k] − x| gives wrong answers when both values are equal but on the same side of x.

Comparing the approaches

Edge cases to test

Variations