leetcode 658. Find K Closest Elements

https://leetcode.com/problems/find-k-closest-elements/

给定一个已排序的数组，以及整数x和整数k，找到数组中最接近x的k个元素存在一个列表中返回。

当有两个元素跟x的差距一样时，优先选择比较小的元素。

一、问题分析

测试用例：

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


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

从第一个例子可以看到，尽管5 - 3 = 3 - 1，不过由于1比5小，所以我们返回的是[1, 2, 3, 4]而不是[2, 3, 4, 5]。

从第二个例子可以看到，此问题存在两种特殊情况：

1）x小于等于数组的最小值，此时只需要把数组前k个元素返回即可；

2）x大于等于数组的最大值，此时只需要把数组后k个元素返回即可

对于其他一般情况，我们用二分法找到第一个大于等于x的元素的index，然后定义low为index左边距离index k+1个单位的位置（不超过0时）或者为0，定义high为index右边距离index k-1个单位的位置（不超过数组的长度时）或者为arr.length - 1。那么，[low, high]这个区间就一定包含k个最接近x的元素。

接着，我们再使用双指针法不断缩小[low, high]的跨度，直到此区间的长度变为k。此时，[low, high]就是我们想找的k个元素。

二、代码实现

class Solution {
    public List<Integer> findClosestElements(int[] arr, int k, int x) {
        
        int start = 0, end = 0;
        
        if (x <= arr[0]) {
            start = 0;
            end = k - 1;
        } else if (x >= arr[arr.length - 1]) {
            start = arr.length - 1 + 1 - k;
            end = arr.length - 1;
        } else {
            //the first one which >= x
            int index = binarySearch(arr, x);

            //[low, high] include the k elements nearest x, its length is >= k
            int low = Math.max(0, index - k - 1);        //int low = Math.max(0, index - k);也可以
            int high = Math.min(arr.length-1, index+k-1);
            
            //high-low+1 is the length of [low, high]
            //when it become k, [low, high] is the k elements nearest x
            while (high - low + 1 > k) {
                if (x - arr[low] <= arr[high] - x) {
                    high --;
                } else if (x - arr[low] > arr[high] - x) {
                    low++;
                } else {
                    if (high - low + 1 - 2 >= k) {
                        low++;
                        high--;
                    } 
                    //差距相等时，优先选择较小的元素
                    if (high - low + 1 - 1 == k) {
                        high--;
                    }
                }     
            }
            
            start = low;
            end = high;
        }
        
        //now [low, high] is the answer
        List<Integer> result = new ArrayList<>();
        while (start <= end) {
            result.add(arr[start]);
            start++;
        }
        
        return result;
    }
    
    public int binarySearch(int[] arr, int target) {
        if (arr == null || arr.length == 0) {
            return -1;
        }
        
        int left = 0, right = arr.length - 1;
        while (left < right) {
            int mid = left + (right - left) / 2;
            
            if (arr[mid] >= target) {
                right = mid;
            } else {
                left = mid + 1;
            }
        }
        
        if (arr[left] >= target) {
            return left;
        }
        return -1;
    }
    
}