Skip to main content

Find K Closest Elements


Find K Closest Elements: 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 = [1,2,3,4,5], k = 4, x = 3
Output: [1,2,3,4]

Example 2:
Input: arr = [1,2,3,4,5], k = 4, x = -1
Output: [1,2,3,4]

Constraints:
  • 1 <= k <= arr.length
  • 1 <= arr.length <= 10^4
  • arr is sorted in ascending order.
  • -10^4 <= arr[i], x <= 10^4

Try this Problem on your own or check similar problems:

  1. Guess Number Higher or Lower
  2. Guess Number Higher or Lower II
  3. Find K-th Smallest Pair Distance
Solution:
public List<Integer> findClosestElements(int[] arr, int k, int x) {
int start = 0, end = arr.length - k;

while(start < end){
int mid = start + (end - start) / 2;
if(x - arr[mid] > arr[mid+k] - x){
start = mid + 1;
}else{
end = mid;
}
}

return Arrays.stream(arr).boxed().toList().subList(start, start + k);
}

Time/Space Complexity:
  • Time Complexity: O(logn)
  • Space Complexity: O(1)

Explanation:

We implement a modified binary search that performs the search in logarithmic time O(log n). The implementation is short and there are only a few things to clarify:

  • Why do we define end = arr.length - k? We do this because we're searching for the start (left most index) of the subarray (start, start + k) containing elements closest to x, so the end can be at most arr.length-k (note that invariant holds start < end, so we will not have an outside boundary index)
  • Now to the decision-making process? How can we be sure that each iteration we're eliminating half of search space? We have two cases:
  1. x - arr[mid] < arr[mid + k] - x we know the x is closer to the arr[mid] and it can be either outside the range mid..mid+k or inside of it but still closer to arr[mid]. To cover k closest elements arr[mid] and arr[mid+k] we have to move to the left end = mid since element at mid is closest element to x that we have so far
  2. x - arr[mid] > arr[mid + k] - x it's the opposite from the first scenario x is closer to arr[mid+k] and it's inside or outside of the range, but either way we have to shift to the right start = mid + 1. mid and all elements before the mid are too far from x so we can skip them and move the left boundary to element after mid

Why don't we use absolute values? When we have A[mid] == A[mid + k] (duplicate can occur in the input array), we can't decide in which direction should we move. This is handled if we don't include the absolute values e.g. arr[mid] = arr[mid+k] = 2, and if we have x = 3, we will have true for x - arr[mid] > arr[mid+k] - x since 3 - 2 > 2 - 3 -> 1 > -1, we know we have to move to the right side. If we have arr[mid] = arr[mid+k] = 2 and x = 1 we will have true for x - arr[mid] < arr[mid+k] - x since 1 - 2 < 2 - 1 -> -1 < 1 and we know we have to move to the left end = mid.