[Leetcode]658.Find K Closest Elements

Posted 2021-03-13 Jamest

链接：LeetCode658

给定一个排序好的数组，两个整数 k 和 x，从数组中找到最靠近 x（两数之差最小）的 k 个数。返回的结果必须要是按升序排好的。如果有两个数与 x 的差值一样，优先选择数值较小的那个数。

示例?1:

输入: ([1,2,3,4,5], k=4, x=3)
输出: ([1,2,3,4])
?

示例 2:

输入: ([1,2,3,4,5], k=4, x=-1)
输出: ([1,2,3,4])

相关标签：二分查找

对于在排序好的数组寻找数，使用二分是常见的思路。这里我们只要找到与x最接近的数，然后在这个数两边找k个即可。
代码如下：

python:

class Solution:
    def findClosestElements(self, arr: List[int], k: int, x: int) -> List[int]:
        ind = self.findColestIndex(arr,x)
        lo = hi = ind
        res = [arr[ind]]
        for i in range(k-1):
            if lo-1<0:
                hi+=1
                res.append(arr[hi])
            elif hi+1>=len(arr):
                lo -= 1
                res.insert(0,arr[lo])
            elif x-arr[lo-1]<=arr[hi+1]-x:
                lo -= 1
                res.insert(0,arr[lo])
            else:
                hi+=1
                res.append(arr[hi])
        # res.sort()
        return res
            

    def findColestIndex(self,arr,x):
        lo = self.findLoIndex(arr,x)
        hi = self.findHiIndex(arr,x)
        if lo == -1:
            return hi 
        if  hi == -1:
            return lo
        return lo if x-arr[lo] <= arr[hi]-x else hi
    
    # 仅小于x的数
    def findLoIndex(self,arr,x):
        lo,hi = 0,len(arr)-1
        while lo<=hi:
            mid = lo+((hi-lo)>>1)
            if x <= arr[mid]:
                hi = mid -1
            else:
                lo = mid +1
        return hi if hi >= 0 else -1

    # 第一个大于等于x的数
    def findHiIndex(self,arr,x):
        lo,hi = 0,len(arr)-1
        while lo<=hi:
            mid = lo+((hi-lo)>>1)
            if x <= arr[mid]:
                hi = mid -1
            else:
                lo = mid +1
        return lo if lo < len(arr) else -1

C++:

class Solution {
public:
    vector<int> findClosestElements(vector<int>& arr, int k, int x) {
        int closest_ind = getClosestIndex(arr,x);
        int lo = closest_ind,hi = closest_ind;
        vector<int> res;
        res.push_back(arr[lo]);
        for(int i=0;i<k-1;i++){
            if(lo<=0){
                res.push_back(arr[++hi]);
            }
            else if(hi>=arr.size()-1){
                res.push_back(arr[--lo]);
            }
            else if(x-arr[lo-1]<=arr[hi+1]-x) {
                res.push_back(arr[--lo]);
            }
            else {
                res.push_back(arr[++hi]);
            }
        }
        sort(res.begin(),res.end());
        return res;
        
    }
    
    int getClosestIndex(vector<int>& arr,int x){
        int hi = getHiIndex(arr,x);
        if(hi==-1){
            return arr.size()-1;
        }
        if(hi == 0){
            return 0;
        }
        int lo = hi-1;
        if(x-arr[lo]<=arr[hi]-x){
            return lo;
        }
        else{
            return hi;
        }
        
    }
    
    
    int getHiIndex(vector<int>& arr,int x){
        int lo=0,hi=arr.size()-1;
        while(lo<=hi){
            int mid = lo+((hi-lo)>>1);
            if (x<=arr[mid]) {
                hi = mid - 1 ;
            }
            else{
                lo = mid + 1;
            }
        }
        if(lo<arr.size()){
            return lo;
        }
        return -1;
    }
};