Find K Closest Elements

Patrick Leaton
Problem Description

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

 

The description was taken from https://leetcode.com/problems/find-k-closest-elements/.

Problem Solution

#O(Log(N)) Time, O(1) Additional Space
class Solution:
    def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
        left = 0
        right = len(arr) - k
       
        while left < right:
            mid = (left+right)//2
            if x-arr[mid] > arr[mid+k]-x:
                left = mid + 1
            else:
                right = mid
       
        return arr[left:left+k]

Problem Explanation


From the name, we may expect this problem can be solved with a heap.  

We could do that, but that would be slower and would cost more space than the optimal solution.

The array is sorted and we are just needing to find the window that contains the k closest elements.  In order to find the correct window, we just need to find the first, closest value within the window and then return the k elements from that. 

That is where a binary search comes into play.

We will compare the closeness of two potential candidates by calculating their difference from the target.  Whichever candidate is closer is where we will begin moving our search towards until our pointers overlap, at which point we will have the beginning of the window.


Let's start by initializing our pointers.

We will initialize our left pointer at the beginning of the array per usual.  However, we will want to initialize our right pointer to the end of the array minus k, since that is the rightmost position that the leftmost value of the window could possibly be in, if we are expected to return k elements.

        left = 0
        right = len(arr) - k

 

Next, let's begin our binary search.

While the left pointer and the right pointer haven't overlapped, we still have values we need to search for the closest target.

        while left < right:

 

At the beginning of each iteration, we will take an average between the two pointers to make a mid pointer.

            mid = (left+right)//2

 

We'll compare the closeness of two values in each iteration to determine which search space we should continue down.

These values are the element at the midpoint and the element k places from there.

If the value at the midpoint is lower but the two values are equal in difference, then the midpoint value wins according to the description.

At that point, we will move the right pointer to the left because our window likely starts leftward in the array.

Otherwise, we will move the left pointer to the right, past the midpoint.

            if x-arr[mid] > arr[mid+k]-x:
                left = mid+1
            else:
                right=mid

 

Once the pointers overlap, we will have the leftmost element of the window, in which case we will need to return the k closest elements from there.

        return arr[left:left+k]

Let's draw this out.  We'll use [2,4,6,8,10] as our test case with x being six and k being three.

 

Our left and right pointers will start here:

[2,4,6,8,10]


So, we have our middle pointer here:

[2,4,6,8,10]

 

We will compare that value, with the one k indices away, to see which one is closest to x.

[2,4,6,8,10]

 

Four is closer.

[2,4,6,8,10]

 

Once more, the left and right pointers are set.

[2,4,6,8,10]

 

Set the mid pointer.

[2,4,6,8,10]

 

We will compare that value, with the one k indices away, to see which one is closest to x.

[2,4,6,8,10]

 

Eight is closer.

[2,4,6,8,10]

 

After we move the left pointer one to the right of mid, we will have our pointers overlap, in which case we have found our k closest elements.

[2,4,6,8,10]