Find K Closest Elements

Given a sorted array, two integerskandx, find thekclosest elements toxin 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:

1. The value k is positive and will always be smaller than the length of the sorted array.

2. Length of the given array is positive and will not exceed 10^4

3. Absolute value of elements in the array and x will not exceed 10^4

UPDATE (2017/9/19): Thearrparameter had been changed to an array of integers(instead of a list of integers).Please reload the code definition to get the latest changes.

Analysis

This is probably one of my favorite binary search problem.

Intuitively, one would find the element x in the array first, and then use two pointers or something else to find the k closest neighbors around the x. However, this problem has a much more elegant approach, which is just using binary search itself.

It's hard to think at first about modifying the condition statements in the binary search, though, without deep understanding of how binary search works.

Instead of finding x position, we can find the left boundary index of the k elements closest to x.

Therefore the condition becomes:

if (x - arr[mid] > arr[mid+k] - x)

And since we prefer the smaller indexes, we exclude the "=" equal in the comparison, thus we move left when the x - arr[mid] = arr[mid+k] - x, to make sure we prioritize the smaller indexes when there's a tie.

Another thing to be careful is setting the initial left and right boundaries (indexes), since we will use arr[mid + k] , thus, we need to make sure the starting right boundaries is less than arr.length - k to avoid index out of bound error.

这题最妙的地方在于

1. 搜索k个距离（数组元素之差的绝对值，而不是下标的距离）最近的数中最小的那个下标，这样通过k就可知道右边界的位置

2. 改变binary search中的条件判断语句，改为k个元素中最小值与最大值与x元素的距离绝对值大小的比较 x - arr[mid] vs arr[mid+k] - x

3. left, right初始值的设定，因为需要[mid + k]，所以right的初始值需要（从arr.length - 1）缩小至 arr.length - k - 1 以防止数组下标越界 （根据Template #1，#3，如果是Template #2 则需从 arr.length 改为 arr.length - k）

Solution

BS Template #1

class Solution {
    public List<Integer> findClosestElements(int[] arr, int k, int x) {
        int left = 0, right = arr.length - k - 1;
        while (left <= right) {
            int mid = left + (right - left) / 2;
            if (x - arr[mid] > arr[mid + k] - x) {
                left = mid + 1;
            } else {
                right = mid - 1;
            }
        }
        List<Integer> ans = new ArrayList<>();
        for (int i = 0; i < k; i++) {
            ans.add(arr[left + i]);
        }
        return ans;
    }
}

BS Template #2

class Solution {
    public List<Integer> findClosestElements(int[] arr, int k, int x) {
        int lo = 0, hi = arr.length - k;
        while (lo < hi) {
            int mid = (lo + hi) / 2;
            if (x - arr[mid] > arr[mid+k] - x)
                lo = mid + 1;
            else
                hi = mid;
        }
        List<Integer> result = new ArrayList<>();
        for(int i=lo;i<lo+k;i++){
            result.add(arr[i]);          
        }
        return result;
    }
}

Reference Solution with comment from LeetCode @meganlee

public List<Integer> findClosestElements(int[] arr, int k, int x) {
    //-------- Main idea behind the binary search algorithm ----------//
    // 1) res will be a consecutive subarray of k size
    // 2) say if we need to pick 4 elems, now we r looking at 5 elem n1, n2, n3, n4, n5
    //    we need to compare two ends: diff(x, n1) and diff(x, n5), 
    //    the number with bigger diff on the end will be eleminated
    //----------------------------------------------------------------//
    // lo and hi: range of all possible start of subarray with size k + 1, so we could compare both ends
    int lo = 0, hi = arr.length - k - 1;
    while (lo <= hi) {
        int mid = lo + (hi - lo) / 2;
        // for subarray starting at mid with size k+1, we compare element of two ends to eliminate the loser
        if (Math.abs(x - arr[mid]) > Math.abs(x - arr[mid+k])) {
            // arr[mid] is the one further away from x, eliminate range[lo, mid]
            lo = mid + 1; 
        } else {
            // arr[mid+k] is the one further away, all [mid, hi] will have simiar situation, elimiate them
            hi = mid - 1; 
        }                
    }     
    // return subarray
    List<Integer> res = new ArrayList(k);
    for (int i = 0; i < k; i++) {
        res.add(arr[lo + i]);
    }
    return res;
}

Reference

https://leetcode.com/problems/find-k-closest-elements/discuss/106419/O(log-n)-Java-1-line-O(log(n)-+-k)-Ruby