658. Find K Closest Elements

Posted 2020-10-02

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)

 1 class Solution {
 2 public:
 3     vector<int> findClosestElements(vector<int>& arr, int k, int x) {
 4         map<int,vector<int>> m;
 5         vector<int> res;
 6         for (int i=0;i<arr.size();i++){
 7             int distance=abs(arr[i]-x);
 8             m[distance].push_back(i);
 9         }
10         int count=0;
11         for (auto it=m.begin();it!=m.end();it++){
12                 int size=it->second.size();
13                 for (int i=0;i<size;i++){
14                     if (count<k){
15                         int index=it->second[i];
16                         res.push_back(arr[index]);
17                         count++;
18                     }                    
19                 }
20                                     
21         }
22         sort(res.begin(),res.end());
23         return res;
24     }
25 };

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).

 1 class Solution {
 2 public:
 3     vector<int> findClosestElements(vector<int>& arr, int k, int x) {
 4         int index=lower_bound(arr.begin(),arr.end(),x)-arr.begin(); //binary search:log(n)
 5         int i=index-1,j=index;
 6         while(k--){
 7             if (i<0||(j<arr.size()&&abs(arr[j]-x)<abs(arr[i]-x))) j++;
 8             else i--;
 9         }
10         vector<int> res(arr.begin()+i+1,arr.begin()+j);
11         return res;
12     }
13 };