【Leetcode】973. K Closest Points to Origin

题目地址：

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

给定一个数组，每个位置是一个平面坐标系下的点。求距离原点欧几里得距离最近的$K$个点。

法1：堆。维护一个大顶堆，先把前$K$个点offer进去，对于接下来的点，如果其离原点比堆顶近，就把堆顶poll出去，让这个点进堆；否则略过。遍历完成后堆里的$K$个点即为所求。

class Solution {
 public:
  vector<vector<int>> kClosest(vector<vector<int>>& ps, int k) {
    auto dis = [](auto& p) { return p[0] * p[0] + p[1] * p[1]; };
    auto cmp = [&](auto& p1, auto& p2) { return dis(p1) < dis(p2); };
    priority_queue<vector<int>, vector<vector<int>>, decltype(cmp)> heap(cmp);
    for (int i = 0; i < k; i++) heap.push(ps[i]);
    for (int i = k; i < ps.size(); i++)
      if (dis(ps[i]) < dis(heap.top())) {
        heap.pop();
        heap.push(ps[i]);
      }

    vector<vector<int>> res;
    for (int i = 0; i < k; i++) {
      res.push_back(heap.top());
      heap.pop();
    }

    return res;
  }
};

时间复杂度$O(n\log k)$，空间$O(k)$。

法2：快速选择（Quick Select）。

class Solution {
 public:
  vector<vector<int>> kClosest(vector<vector<int>>& ps, int k) {
    auto dis = [](auto& p) { return p[0] * p[0] + p[1] * p[1]; };
    quick_select(ps, 0, ps.size() - 1, k - 1, dis);
    return vector<vector<int>>(ps.begin(), ps.begin() + k);
  }

  void quick_select(vector<vector<int>>& ps, int l, int r, int rk, auto& dis) {
    if (l >= r) return;
    int i = l, j = r;
    int piv_dis = dis(ps[l + (r - l >> 1)]);
    while (i <= j) {
      while (dis(ps[i]) < piv_dis) i++;
      while (dis(ps[j]) > piv_dis) j--;
      if (i <= j) swap(ps[i++], ps[j--]);
    }
    if (rk <= j) quick_select(ps, l, j, rk, dis);
    if (rk >= i) quick_select(ps, i, r, rk, dis);
  }
};

平均时间复杂度$O(n)$，空间复杂度$O(\log n)$。
算法正确性可参照快速排序的证明https://blog.csdn.net/qq_46105170/article/details/104085776，算法复杂度可参照https://blog.csdn.net/qq_46105170/article/details/104041143。