【堆】B010_LC_最接近原點的 K 個點（3 行排序 / 快速選擇）

一、Problem

我們有一個由平面上的點組成的列表 points。需要從中找出 K 個距離原點 (0, 0) 最近的點。

（這裏，平面上兩點之間的距離是歐幾里德距離。）

你可以按任何順序返回答案。除了點座標的順序之外，答案確保是唯一的。

輸入：points = [[1,3],[-2,2]], K = 1
輸出：[[-2,2]]
解釋： 
(1, 3) 和原點之間的距離爲 sqrt(10)，
(-2, 2) 和原點之間的距離爲 sqrt(8)，
由於 sqrt(8) < sqrt(10)，(-2, 2) 離原點更近。
我們只需要距離原點最近的 K = 1 個點，所以答案就是 [[-2,2]]。

提示：

1 <= K <= points.length <= 10000
-10000 < points[i][0] < 10000
-10000 < points[i][1] < 10000

二、Solution

方法一：排序

…

class Solution {
public:
	typedef pair<int, int> pii;
	int dist(vector<int>& point) {
        return point[0]*point[0] + point[1]*point[1];
    }
    vector<vector<int>> kClosest(vector<vector<int>>& ps, int K) {
    	int n = ps.size();
    	vector<vector<int>> ans;
    	vector<pii> vt;
    	for (int i = 0; i < n; i++) {
    		vt.push_back({dist(ps[i]), i});
    	}

    	sort(vt.begin(), vt.end(), [&](pii a, pii b) {
    		return a.first < b.first;
    	});

    	for (int i = 0; i < K; i++) {
    		int id = vt[i].second;
    		ans.push_back(ps[id]);
    	}
    	return ans;
    }
};

我們還可以將空間壓縮到 O(1)

class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>>& ps, int K) {
    	sort(ps.begin(), ps.end(), [&](auto a, auto b) {
    		return a[0]*a[0]+a[1]*a[1] < b[0]*b[0]+b[1]*b[1];
    	});
        ps.resize(K);
    	return ps;
    }
};

複雜度分析

• 時間複雜度：$O(nlogn)$，
• 空間複雜度：$O(1)$，

---

方法二：快速選擇

我們不需要這 K 個點完全有序，我們只需要知道這 K 個元素是最小的 K 個即可，所以我們可以用 quickSelect 做

從左往右找到第一個 dist(i) >= x 的下標 i，從右往左找第一個 <= x 的下標 j，將他們 ps[i] 和 ps[j] 交換，反覆如此，直到 i >= j

class Solution {
public:
	vector<vector<int>> ps;
	int dist(vector<int>& point) { return point[0]*point[0] + point[1]*point[1]; }
    void quick_select(int l, int r, int k) {
        if (l >= r)
            return;
    	int i = l-1, j = r+1, p = rand()%(r-l+1)+l, x = dist(ps[p]);
        
        while (i < j) {
            do i++; while (dist(ps[i]) < x);
            do j--; while (dist(ps[j]) > x);
            if (i < j) 	   swap(ps[i], ps[j]);
        }
    	if (j >= k) quick_select(l, j, k);
    	else 		quick_select(j+1, r, k);
    }
    vector<vector<int>> kClosest(vector<vector<int>>& points, int K) {
    	ps = points;
    	quick_select(0, ps.size()-1, K);
    	ps.resize(K);
    	return ps;
    }
};

複雜度分析

• 時間複雜度：$O(n)$，
• 空間複雜度：$O(1)$，