658. Find K Closest Elements

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:

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⁴
3. Absolute value of elements in the array and x will not exceed 10⁴

Solution 1: set a map, the key is |arr[i]-x|, the value is the index vector, eg, 1:[1,3] for example 1 input.  O(klogk+n)

class Solution {
public:
    vector<int> findClosestElements(vector<int>& arr, int k, int x) {
        map<int,vector<int>> m;
        vector<int> res;
        for (int i=0;i<arr.size();i++){
            int distance=abs(arr[i]-x);
            m[distance].push_back(i);
        }
        int count=0;
        for (auto it=m.begin();it!=m.end();it++){
                int size=it->second.size();
                for (int i=0;i<size;i++){
                    if (count<k){
                        int index=it->second[i];
                        res.push_back(arr[index]);
                        count++;
                    }                    
                }
                                    
        }
        sort(res.begin(),res.end());
        return res;
    }
};

Solution 2: use the binary search to find the first number which is >=x, then, determine the indices of the start and the end of a subarray in arr, where the subarray is our result. The time complexity is O(logn + k).

class Solution {
public:
    vector<int> findClosestElements(vector<int>& arr, int k, int x) {
        int index=lower_bound(arr.begin(),arr.end(),x)-arr.begin(); //binary search:log(n)
        int i=index-1,j=index;
        while(k--){
            if (i<0||(j<arr.size()&&abs(arr[j]-x)<abs(arr[i]-x))) j++;
            else i--;
        }
        vector<int> res(arr.begin()+i+1,arr.begin()+j);
        return res;
    }
};

 line 4 use the lower_bound which is a binary search function like below.

int binarySearch(vector<int>& arr,int x){
        int beg=0,end=arr.size()-1,pos;
        while (beg<end){
            int mid=(beg+end)/2;
            if (arr[mid]==x) return mid;
            if (arr[mid]<x) {beg=mid+1;pos=beg;}
            if (arr[mid]>x) {end=mid;pos=end;}
        }
        return pos;
    }