658

658. Find K Closest Elements

You can solve this question with a heap, however, because the array is already sorted, an optimal solution would use binary search. We’re looking for the array of length k that has elements as close as possible to input x. Closeness is defined as the difference between arr[num] - x.

During binary search, if the middle element is less closer than the element at m + k, we know that all elements before the middle element can’t be in the answer (because the array is sorted). We move the left pointer to m + 1. Otherwise, we move the right pointer to m.

Eventually the left pointer will be at the start of the answer subarray.

The solution is as follows:

class Solution:
    def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
        l, r = 0, len(arr) - k
 
        while l < r:
            m = (l + r) // 2
 
            if x - arr[m] > arr[m + k] - x:
                l = m + 1
            else:
                r = m
 
        return arr[l:l+k]

_ Time Complexity:

O(log(n - k) + k) - The binary search takes log(n - k) and building the answer array takes k time.

_ Space Complexity:

O(1) - We don’t use any extra space.