Tags: "leetcode", "binary-search", access_time 2-min read

Edit this post on Github

Find K Closest Elements

Created: March 14, 2019 by [lek-tin]

Last updated: March 14, 2019

Given a sorted array, two integers k and x, find the k closest elements to x in the array. The result should also be sorted in ascending order. If there is a tie, the smaller elements are always preferred.

Example 1

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

Example 2

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

Note

  • The value k is positive and will always be smaller than the length of the sorted array.
  • Length of the given array is positive and will not exceed 104
  • Absolute value of elements in the array and x will not exceed 104

Solution:

Assume we are taking arr[i] ~ arr[i + k -1]. We can binary search i. We compare the distance between x - arr[mid] and arr[mid - k] - x. If x - arr[mid] > arr[mid + k] - x, it means arr[mid + 1] ~ A[mid + k] is better than arr[mid] ~ arr[mid + k - 1], and we have mid smaller than the right i. So assign left = mid + 1.

// Time: O(log(N-K) + K)
class Solution {
    public List<Integer> findClosestElements(int[] arr, int k, int x) {
        Integer low = 0, high = arr.length - k;
        List<Integer> res = new ArrayList<>();

        while (low < high) {
            int mid = (low + high) / 2;
            if (x - arr[mid] > arr[mid + k] - x) {
                low = mid + 1;
            } else {
                high = mid;
            }
        }

        for (int i = low; i < low + k; i++) res.add(arr[i]);

        return res;
    }
}
class Solution:
    def findClosestElements(self, arr, k, x):
        """
        :type arr: List[int]
        :type k: int
        :type x: int
        :rtype: List[int]
        """
        low = 0
        high = len(arr) - k

        while low < high:
            mid = (low + high) // 2
            if x > (arr[mid] + arr[mid + k]) / 2:
                low = mid + 1
            else:
                high = mid

        return arr[low:low + k]