K Closest Points

https://leetcode.com/problems/k-closest-points-to-origin/

从一堆点中找离原点欧式距离最近的K个点。找前K用quick select或heap。找最近，用最大堆。

Code

/*
 * @lc app=leetcode id=973 lang=cpp
 *
 * [973] K Closest Points to Origin
 */

// @lc code=start
class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int K) {
        const auto dist = [](vector<int> &p) {
            return p[0] * p[0] + p[1] * p[1];
        };
        int l = 0, r = points.size() - 1;
        const auto partition = [&](int l, int r){
            const auto t = rand() % (r - l + 1) + l;
            const int s = l;
            swap(points[l++], points[t]);
            const auto td = dist(points[s]);
            while (l <= r) {
                const auto ld = dist(points[l]);
                const auto rd = dist(points[r]);
                if (ld <= td) {
                    ++l;
                } else if (rd > td) {
                    --r;
                } else {
                    swap(points[l], points[r]);
                    ++l;
                    --r;
                }
            }
            swap(points[s], points[r]);
            return r;
        };
        while (true) {
            int pos = partition(l, r);
            if (pos == K - 1) break;
            else if (pos < K - 1) {
                l = pos + 1;
            } else {
                r = pos - 1;
            }
        }
        vector<vector<int>> res(K);
        for (int i = 0; i < K; ++i) {
            res[i] = points[i];
        }
        return res;
    }
};
// @lc code=end

Analysis

平均O(N), 最差O(N^2).