658. Find K Closest Elements

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

---

Thoughts

• A Straight-forward Binary Search Solution
  We can use binary search template to find the first number which is larger or equals to x, (say pivot).
  Then from pivot position and pivot - 1, we use the 2-pointers to find the closest numbers until we find k numbers.

  • Time complexity: O(logn): find the pivot + O(k):find the k numbers + O(klogk): sort the result array
  • Space complexity: O(k) for result array

• Binary Search to find the left bound of the k window.
  The answer should be a k-length window. To get the k-length window we only need to find the start position (left bound of the k- window).
  We can use binary search to find the left bound. But we need to know how to define if the mid number is the start position or not.
  Think it this way: from position mid, we look at the k numbers after mid. So it’s arr[mid: mid + k + 1]. This window has k + 1 numbers in it. And either mid or mid + k is the one that we need to delete.

  • If x - arr[mid] > arr[mid + k] - x ⇒ for sure arr[mid] is not in the k numbers and all the numbers before mid should not be the start.
  • If x - arr[mid] <= arr[mid + k] - x ⇒ for sure arr[mid + k] is not in the k numbers and all the numbers after mid + k should not be the result. The left bound should be mid or someone before mid.
  • Time complexity: O(logn): find the left bound + O(k):find the k numbers
  • Space complexity: O(k) for result array

Code

leverage a regular binary search template

    def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
        pivot = self.find_pivot(arr, x)
        result = self.find_k_neighbors(arr, pivot, k, x)
        result.sort()
        return result
            
    '''
    Find the first number which is larger or equals to x
    Binary Search Template
    '''
    def find_pivot(self, arr, x):
        start, end = 0, len(arr) - 1
        while start < end - 1:
            mid = (start + end) // 2
            if arr[mid] >= x:
                end = mid
            else:
                start = mid                
        if arr[start] >= x:
            return start
        else:
            return end
    
        self.find_k_neighbors(arr, pivot, result, k, x)
        result.sort()
        return result
            
    def find_k_neighbors(self, arr, pivot, k, x):
        result = []
        left, right = pivot - 1, pivot
        '''
        Only need to find the k closest number:
        for loop k times
        '''
        for _ in range(k):
            if self.is_left_closer(left, right, arr, x):
                    result.append(arr[left])
                    left -= 1
            else:
                    result.append(arr[right])
                    right += 1
        return result
       
    '''
    When finding the k numbers, we need to check which one has reach the end and continue adding numbers to result from the other side until we find k numbers. 
    It's easier to define a compare funtion here to solve the problem rather than check after the while loop.
    while 0 <= left and right <= len(arr) - 1:
          if x - arr[left] <= arr[right] - x:
              result.append(arr[left])
              left -= 1
          else:
              result.append(arr[right])
              right += 1
          if len(result) == k:
              return
    if left < 0 and right <= len(arr) - 1:
    	while len(result) < k:
            result.append(arr[right])
            right += 1
    else:
        while len(result) < k:
            result.append(arr[left])
            left -= 1
    '''
    def is_left_closer(self, left, right, arr, x):
        if left < 0:
            return False
        if right >= len(arr):
            return True
        return x - arr[left] <= arr[right] - x

a smart binary search solution

    def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
    	'''
        Find the left bound of the k window.
        Since we need to find k numbers, so the left bound can be found in [0, n - k]
    	'''
        left, right = 0, len(arr) - k
        
        while left < right:
            mid = (left + right) // 2
            '''
            For the k + 1 length array [mid, mid + k] (both sides included),
            find which side should not included in this k array.
            If x - a and b - x are the same and a < b, then a should be included.
            '''
            if x - arr[mid] > arr[mid + k] - x:
            '''
            arr[mid + k] is closer, so mid should not be in the result array.
            The start position should be anyone who is large than mid.
            '''
                left = mid + 1
            else:
            '''
            arr[mid] is closer, so mid + k should not be in the result array.
            The start position should be anyone who is before or equals to mid.
            '''
                right = mid
        '''
        return the k values from left
        This will also cost O(k)
        '''
        return arr[left: left + k]

---

Summary

We can use binary search if we want to find some index or value which is related to a target value in the ordered array.

• find the any / first / last position of target
• find the times that target number appears: last position - first position + 1
• find the first number / position which equals or is larger / smaller than target
• find where to insert target number: the first position which is smaller / larger than target

---