973. K Closest Points to Origin

這題我一開始就很單純的算距離, 比大小, 結果可以AC但速度太慢了

O(nlogn)

class Solution {
public:
    static bool compareClosest(vector<int> a, vector<int> b){
        double a_dis = sqrt(a[0] * a[0] + a[1] * a[1]);
        double b_dis = sqrt(b[0] * b[0] + b[1] * b[1]);
        return a_dis < b_dis;
    }
    vector<vector<int>> kClosest(vector<vector<int>>& points, int k) {
        sort(points.begin(), points.end(), compareClosest);
        while(points.size() != k){
            points.pop_back();
        }
        return points;
    }
};

但這可以用priority_queue做, 可以縮到O(nlogk)

class Solution {
public:
    using distanceMapping = pair<int, int>;
 
    vector<vector<int>> kClosest(vector<vector<int>>& points, int k) {
        priority_queue<distanceMapping> q;
        for(int i = 0; i < points.size();i++){
            auto p = points[i];
            int dis = p[0] * p[0] + p[1] * p[1];
            if(q.empty() || q.size() < k){
                q.push({dis, i});
            }else{
                auto cur = q.top();
                if(cur.first > dis){
                    q.pop();
                    q.push({dis, i});
                }
            }
        }
        vector<vector<int>> res;
        while(!q.empty()){
            auto cur = q.top();
            res.push_back(points[cur.second]);
            q.pop();
        }
        return res;
    }
};

會快很多, 還有一種better case可以接近O(n)的, 使用nth_element, 這個函式幫忙

class Solution {
public:
    using distanceMapping = pair<int, int>;
 
    vector<vector<int>> kClosest(vector<vector<int>>& points, int k) {
        nth_element(points.begin(), points.begin() + k - 1, points.end(),[](vector<int> &a, vector<int> &b){
            return a[0] * a[0] + a[1] * a[1] < b[0] * b[0] + b[1] * b[1];
        });
        
        return vector<vector<int>>(points.begin(), points.begin() + k);
    }
};