leetcode 658. Find K Closest Elements

淩亂°似流年 2023-06-21 15:24 42阅读 0赞

[https://leetcode.com/problems/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;
        }
        
    }

[https_leetcode.com_problems_find-k-closest-elements]: https://leetcode.com/problems/find-k-closest-elements/