Leetcode 658: Find K Closest Elements

The original question can be found here.

The question is asking to return the k closet elements to target x. The array is sorted.

Heap

First solution is using a heap to keep track the current closest elements to x. Of course the comparator needs to be tweaked accordingly.

class Solution {
    public List<Integer> findClosestElements(int[] arr, int k, int x) {
        PriorityQueue<Integer> pq = new PriorityQueue<Integer>((a, b) -> {
            if (Math.abs(a - x) == Math.abs(b - x)) {
                return a >= b ? -1 : 1;
            }
            return Math.abs(a - x) < Math.abs(b - x) ? 1 : -1;
        });
        
        for (int num : arr) {
            pq.add(num);
            if (pq.size() > k) {
                pq.poll();
            }
        }

        List<Integer> res = new ArrayList();
        while (pq.size() > 0) {
            res.add(pq.poll());
        }

        Collections.sort(res);
        return res;
    }
}

We’re using a “max” heap, and the element at root is the one farthest from x. Time complexity: O(nlogk + klogk). Space complexity: O(k). However, this solution does not use leverage the characteristic that array is already sorted.

Binary Search

First we locate the target element in the array, and then we expand to both left and right to find the K closest elements. There’s one case needs to take into consideration: what if we cannot find the element? Actually we don’t have to. We find the first element that is >= target and expand from there (of course you can choose the last element that is <= target which basically the same idea).

class Solution:
    def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
        right = self.find_first_greater_or_equal_to_target_element(arr, x)
        left = right - 1
        res = []
        while len(res) < k:
            if self.is_left_closer(arr, left, right, x):
                res.append(arr[left])
                left -= 1
            else:
                res.append(arr[right])
                right += 1
        res.sort()
        return res
        
    def find_first_greater_or_equal_to_target_element(self, arr: List[int], target: int) -> int:
        left, right = 0, len(arr) - 1
        while left + 1 < right:
            mid = (left + right) // 2
            if arr[mid] < target:
                left = mid
            else:
                right = mid
        if arr[right] >= target:
            return right
        if arr[left] >= target:
            return left
        return len(arr)

    def is_left_closer(self, arr: List[int], left: int, right: int, target: int) -> bool:
        if left < 0:
            return False

        if right >= len(arr):
            return True
        
        return abs(arr[left] - target) <= abs(arr[right] - target)

Time Complexity: O(logn + k). Space Complexity: O(1).

Scroll to Top