Leetcode 658. Find K Closest Elements

Problem

Leetcode 658 - Find K Closest Elements

Example:

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

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

Approach

Since we need to return k elements taht are closest to x, and the input array is sorted. We can use binary search to find the window that is closest to x.

Our initial left = 0, while right = len(arr) - k because we need extra k space for the return array

If arr[mid] is farther from target than arr[mid+k] which is k places ahead of mid then we need to pull left to mid with 1 offset; otherwise we can pull right at mid

Visualization of the Approach:

arr = [1,2,3,4,5], k = 2, x = 5
       l     r
mid = 1, nums[mid] = 2
x - arr[mid] > arr[mid+k] - x -> 5 - 1 > 4 - 5 -> 4 > -1. Thus left = mid+1

arr = [1,2,3,4,5]
           l r
mid = 2, nums[mid] = 3
x - arr[mid] > arr[mid+k] - x -> 5 - 3 > 5 - 5 -> 2 > 0. Thus left = mid+1

arr = [1,2,3,4,5]
             lr
left == right stop while loop

arr[left: left+k] = [4,5]

Here is the Python code for the solution:

class Solution:
    def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
        left, right = 0, 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]

Time Complexity and Space Complexity

Time Complexity: $O(log(n-k)+k)$

Space Complexity: $O(k)$