題目網址：https://leetcode.cn/problems/k-closest-points-to-origin/

題意：給一 array points 和一整數 k, 其中 points[i] = [x_i, y_i] 表示平面上的一個點, 返回距離 (0, 0) 最近的 k 個點。

平面上兩點的距離是採用歐式距離 $\sqrt{(x1 - x2)^2 + (y1 - y2)^2}$

可以按任何順序返回答案。

Solution：

想法：利用 Heap, 用 max heap 維護前 k 小的元素, 類似 347. Top K Frequent Elements

  • 先將前 k 個元素加入到 max heap 中
  • 後面 n - k 個點, 都是先 insert 到 heap 中再 pop 掉, 來維持 heap 中 k 個元素
  • 此時, heap 中為前 k 小的元素, 將其取出即可
typedef tuple<int, int, int> t3i; // {dist, x, y}

class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int k) {
        priority_queue<t3i> pq; // max heap 維護前 k 小

        for (const auto& point : points) {
            const int x = point[0], y = point[1];
            const int dist = x * x + y * y;
            pq.emplace(t3i{dist, x, y});

            // 若加入後 heap 中有 k + 1 個, 則 pop 掉 dist 最大的來維持 k 個
            if (pq.size() > k) {
                pq.pop();
            }
        }

        // 此時, heap 中為前 k 小元素, 將其取出即可
        vector<vector<int>> res;
        while (!pq.empty()) {
            const auto [dist, x, y] = pq.top();
            pq.pop();
            res.emplace_back(vector<int>{x, y});
        }
        return res;
    }
};
  • time：$O(n \cdot log(k))$ ➔ 前面兩個 for loop 最多 insert/delete n 個點, 每次操作需 $O(log(k))$, 因為 pq 中最多 k 個點
  • space：$O(k)$ ➔ pq 最多有 k 個元素