LeetCode第 973 题：最接近原点的 K 个点(C++)

973. 最接近原点的 K 个点 - 力扣（LeetCode）

优先级队列

topK问题，优先级队列的典型应用：

class Solution {
public:
    struct cmp{
        bool operator()(const pair<int, int> &a, const pair<int, int> &b){
            return a.first*a.first + a.second*a.second < b.first*b.first + b.second*b.second;
        }
    };
    vector<vector<int>> kClosest(vector<vector<int>>& points, int K) {
        vector<vector<int>> res;
        priority_queue<pair<int, int>, vector<pair<int, int>>, cmp> q;//大顶堆
        for(int i = 0; i < K; ++i)  q.push({points[i][0], points[i][1]});
        for(int i = K; i < points.size(); ++i){
            if(points[i][0]*points[i][0] + points[i][1]*points[i][1] >= q.top().first*q.top().first + q.top().second*q.top().second)
                continue;
            q.pop();
            q.push({points[i][0], points[i][1]});
        }
        while(!q.empty()){
            res.push_back({q.top().first, q.top().second});
            q.pop();
        }
        return res;
    }
};

优化一下：

class Solution {
public:
    struct cmp{
        bool operator()(const pair<int, int> &a, const pair<int, int> &b){
            return a.first < b.first;
        }
    };
    vector<vector<int>> kClosest(vector<vector<int>>& points, int K) {
        vector<vector<int>> res;
        //默认大顶堆，<距离，下标>，默认按照第一项排序
        priority_queue<pair<int, int>> q;
        //priority_queue<pair<int, int>, vector<pair<int, int>>, cmp> q;//大顶堆
        for(int i = 0; i < K; ++i)  q.push({points[i][0]*points[i][0] + points[i][1]*points[i][1], i});
        for(int i = K; i < points.size(); ++i){
            int val = points[i][0]*points[i][0] + points[i][1]*points[i][1];
            if(val >= q.top().first) continue;
            q.pop();
            q.push({val, i});
        }
        while(!q.empty()){
            res.push_back(points[q.top().second]);
            q.pop();
        }
        return res;
    }
};

变体快排

另外，topK问题还有一种常用解法是partition，就是快速排序的那个partition:数据结构与算法：12 | 排序（下）：归并排序与快速排序_zj-CSDN博客的解答开篇那个部分有提到。

class Solution {
public:
    int partition(vector<vector<int>> &points, int l, int r){
        int pivot = points[r][0]*points[r][0] + points[r][1]*points[r][1];
        int idx = l;
        for(int i = l; i < r; ++i){
            if(points[i][0]*points[i][0] + points[i][1]*points[i][1] < pivot)
                swap(points[idx++], points[i]);
        }
        swap(points[idx], points[r]);
        return idx;
    }

    void random_seltct(vector<vector<int>> &points, int l, int r, int k){//递归直到idx==k-1
        int i = rand()%(r-l+1) + l;
        swap(points[i], points[r]);//随机交换，避免划分最坏情况的出现
        int idx = partition(points, l, r);
        if(idx > k-1 )    random_seltct(points, l, idx-1, k);
        else if(idx < k-1)    random_seltct(points, idx+1, r, k);
    }

    vector<vector<int>> kClosest(vector<vector<int>>& points, int K) {
        random_seltct(points, 0, points.size()-1, K);
        return {points.begin(), points.begin()+K};
    }
};

这一题使用划分的方法效率会相比优先级队列更好。

STL nth_element()

还有更方便的，STL库函数大法:nth_element - C++ Reference，其实原理就是上面的划分思路

简单调用方法：

nth_element(first,nth,last)

nth表示第n大的元素应该存放的位置：

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