我们有一个由平面上的点组成的列表 points。需要从中找出 K 个距离原点 (0, 0) 最近的点。
（这里，平面上两点之间的距离是欧几里德距离。）
你可以按任何顺序返回答案。除了点坐标的顺序之外，答案确保是唯一的。
示例 1：
输入：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]]。
示例 2：
输入：points = [[3,3],[5,-1],[-2,4]], K = 2
输出：[[3,3],[-2,4]]
（答案 [[-2,4],[3,3]] 也会被接受。）
提示：
1 <= K <= points.length <= 10000
-10000 < points[i][0] < 10000
-10000 < points[i][1] < 10000
来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/k-closest-points-to-origin
方法一：排序
思路
将所有点到原点的距离进行排序，然后输出距离最近的 K 个点。
算法
有两种版本的代码。
使用 JAVA 语言，我们创建一个距离数组，然后排序找到第 K 大的距离，然后我们返回距离小于等于这个第 K 大距离的 K 个点。
使用 Python 语言，我们通过自定义一个排序的键值比较函数来完成排序，然后我们返回列表中的前 K 个点。
java
class Solution {
public int[][] kClosest(int[][] points, int K) {
int N = points.length;
int[] dists = new int[N];
for (int i = 0; i < N; ++i)
dists[i] = dist(points[i]);
Arrays.sort(dists);
int distK = dists[K-1];
int[][] ans = new int[K][2];
int t = 0;
for (int i = 0; i < N; ++i)
if (dist(points[i]) <= distK)
ans[t++] = points[i];
return ans;
}
public int dist(int[] point) {
return point[0] * point[0] + point[1] * point[1];
}
}
python
class Solution(object):
def kClosest(self, points, K):
points.sort(key = lambda P: P[0]**2 + P[1]**2)
return points[:K]
c++
class Solution {
public:
vector<vector<int>> kClosest(vector<vector<int>>& points, int K) {
map<vector<int>, double> dist;
for(size_t i=0;i<points.size();++i){
dist[points[i]] = pow(points[i][0],2) + pow(points[i][1], 2);
}
vector<pair<vector<int>, double>> p_dist(dist.begin(), dist.end());
std::sort(p_dist.begin(), p_dist.end(), [](const pair<vector<int>, double>& l, const pair<vector<int>, double>& r){
return l.second < r.second;
});
vector<vector<int>> p_kClosest;
for(size_t i=0;i<static_cast<size_t>(K);++i){
p_kClosest.push_back(p_dist[i].first);
}
return p_kClosest;
}
};
复杂度分析
时间复杂度：O(NlogN)，其中 N 是给定点的数量。
空间复杂度：O(N)。
方法二：分治法
思路
我们想要一个复杂度比 NlogN 更低的算法。 显然，做到这件事情的唯一办法就是利用题目中可以按照任何顺序返回 K 个点的条件，否则的话，必要的排序将会花费我们 NlogN 的时间。
我们随机地选择一个元素 x = A[i] 然后将数组分为两部分： 一部分是到原点距离小于 x 的，另一部分是到原点距离大于等于 x 的。 这个快速选择的过程与快速排序中选择一个关键元素将数组分为两部分的过程类似。
如果我们快速选择一些关键元素，那么每次就可以将问题规模缩减为原来的一半，平均下来时间复杂度就是线性的。
算法
我们定义一个函数 work(i, j, K)，它的功能是部分排序 (points[i], points[i+1], ..., points[j]) 使得最小的 K 个元素出现在数组的首部，也就是 (i, i+1, ..., i+K-1)。
首先，我们从数组中选择一个随机的元素作为关键元素，然后使用这个元素将数组分为上述的两部分。为了能使用线性时间的完成这件事，我们需要两个指针 i 与 j，然后将它们移动到放错了位置元素的地方，然后交换这些元素。
然后，我们就有了两个部分 [oi, i] 与 [i+1, oj]，其中 (oi, oj) 是原来调用 work(i, j, K) 时候 (i, j) 的值。假设第一部分有 10 个元，第二部分有15 个元素。如果 K = 5 的话，我们只需要对第一部分调用 work(oi, i, 5)。否则的话，假如说 K = 17，那么第一部分的 10 个元素应该都需要被选择，我们只需要对第二部分调用 work(i+1, oj, 7) 就行了。
java
import java.util.concurrent.ThreadLocalRandom;
class Solution {
int[][] points;
public int[][] kClosest(int[][] points, int K) {
this.points = points;
work(0, points.length - 1, K);
return Arrays.copyOfRange(points, 0, K);
}
public void work(int i, int j, int K) {
if (i >= j) return;
int oi = i, oj = j;
int pivot = dist(ThreadLocalRandom.current().nextInt(i, j));
while (i < j) {
while (i < j && dist(i) < pivot) i++;
while (i < j && dist(j) > pivot) j--;
swap(i, j);
}
if (K <= i - oi + 1)
work(oi, i, K);
else
work(i+1, oj, K - (i - oi + 1));
}
public int dist(int i) {
return points[i][0] * points[i][0] + points[i][1] * points[i][1];
}
}
python
class Solution(object):
def kClosest(self, points, K):
dist = lambda i: points[i][0]**2 + points[i][1]**2
def work(i, j, K):
if i >= j: return
oi, oj = i, j
pivot = dist(random.randint(i, j))
while i < j:
while i < j and dist(i) < pivot: i += 1
while i < j and dist(j) > pivot: j -= 1
if dist(i) == pivot and dist(j) == pivot: break
points[i], points[j] = points[j], points[i]
if K <= i - oi + 1:
work(oi, i, K)
else:
work(i+1, oj, K - (i - oi + 1))
work(0, len(points) - 1, K)
return points[:K]
复杂度分析
时间复杂度：O(N) ，这是在平均情况下 的时间复杂度， 其中 N 是给定点的数量。
空间复杂度：O(N)。
方法三：小根堆
c++
class Solution {
public:
struct point{
int x,y;
bool operator < (const point& res) const
{
return sqrt(x*x + y*y) > sqrt(res.x*res.x + res.y*res.y);
}
};
vector<vector<int>> kClosest(vector<vector<int>>& points, int K) {
priority_queue<point> q;
for(int i = 0; i < points.size(); i++){
point p;
p.x = points[i][0];
p.y = points[i][1];
q.push(p);
}
vector<vector<int>> res;
for(int i = 0; i < K; i++){
res.push_back(vector<int>());
res[i].push_back(q.top().x);
res[i].push_back(q.top().y);
q.pop();
}
return res;
}
};
方法四：大根堆
java
class Solution {
public int[][] kClosest(int[][] points, int K) {
//大根堆
PriorityQueue<int[]>queue = new PriorityQueue<int[]>(
(o1,o2)->{
return -(o1[0]*o1[0]+o1[1]*o1[1]-o2[0]*o2[0]-o2[1]*o2[1]);
}
);
for(int i = 0;i < K;i++)
queue.offer(points[i]);
for(int i = K;i < points.length;i++)
{
int[]p = queue.peek();
int[]q = points[i];
if(p[0]*p[0]+p[1]*p[1] > q[0]*q[0] + q[1]*q[1])
{
queue.poll();
queue.offer(q);
}
}
int[][]ans = new int[K][2];
Iterator<int[]>iterator = queue.iterator();
int i = 0;
while(iterator.hasNext())
ans[i++] = iterator.next();
return ans;
}
}
c++
class Solution {
struct cmp{
bool operator()(pair<int,int>&a,pair<int,int>&b){
return a.first*a.first+a.second*a.second<b.first*b.first+b.second*b.second;
}
};
public:
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 < points.size(); i++){
pair<int,int> x;
x.first = points[i][0];
x.second = points[i][1];
q.push(x);
if(q.size() > K) q.pop();
}
while(!q.empty()){
res.push_back({q.top().first, q.top().second});
q.pop();
}
return res;
}
};