Find the K closest points to the origin in a 2D plane, given an array containing N points.
/*
public class Point {
public int x;
public int y;
public Point(int x, int y) {
this.x = x;
this.y = y;
}
}
*/
public List<Point> findKClosest(Point[] p, int k) {
PriorityQueue<Point> pq = new PriorityQueue<>(10, new Comparator<Point>() {
@Override
public int compare(Point a, Point b) {
return (b.x * b.x + b.y * b.y) - (a.x * a.x + a.y * a.y);
}
});
for (int i = 0; i < p.length; i++) {
if (i < k)
pq.offer(p[i]);
else {
Point temp = pq.peek();
if ((p[i].x * p[i].x + p[i].y * p[i].y) - (temp.x * temp.x + temp.y * temp.y) < 0) {
pq.poll();
pq.offer(p[i]);
}
}
}
List<Point> x = new ArrayList<>();
while (!pq.isEmpty())
x.add(pq.poll());
return x;
}
这道题和找第k大或第k小的题目的思路基本相同，就是在遍历所有Point的同时，维护一个size为k的max—heap，一旦发现size为k+1，我们就把max-heap头上最大的元素移出heap，因为这里的heap是max-heap，所以heap头部的元素比heap里其他的元素都要比heap里的其他元素离原点远。这样使得heap里的元素是到目前为止里原点最近的k的点。
#include <iostream>
tagsstart
<i class="fa fa-tags" aria-hidden="true"></i> [Linkedin](/tags.html#linkedin) [Square](/tags.html#square)
tagsstop
#include <string>
#include <algorithm>
#include <queue>
#include <math.h>
#include <vector>
using namespace std;
struct Point {
double x;
double y;
Point(double a, double b) {
x = a;
y = b;
}
};
double getDistance(Point a, Point b) {
return sqrt((a.x - b.x) * (a.x - b.x) + (a.y - b.y) * (a.y - b.y));
}
typedef bool (*comp)(Point, Point);
Point global_origin = Point(0,0);
bool compare(Point a, Point b)
{
return (getDistance(a, global_origin)< getDistance(b, global_origin));
}
vector<Point> Solution(vector<Point> &array, Point origin, int k) {
global_origin = Point(origin.x, origin.y);
priority_queue<Point, std::vector<Point>, comp> pq(compare);
vector<Point> ret;
for (int i = 0; i < array.size(); i++) {
Point p = array[i];
pq.push(p);
if (pq.size() > k)
pq.pop();
}
int index = 0;
while (!pq.empty()){
Point p = pq.top();
ret.push_back(p);
pq.pop();
}
return ret;
}
int main()
{
Point p1 = Point(4.5, 6.0);
Point p2 = Point(4.0, 7.0);
Point p3 = Point(4.0, 4.0);
Point p4 = Point(2.0, 5.0);
Point p5 = Point(1.0, 1.0);
vector<Point> array = {p1, p2, p3, p4, p5};
int k = 2;
Point origin = Point(0.0, 0.0);
vector<Point> ans = Solution(array, origin, k);
for (int i = 0; i < ans.size(); i++) {
cout << i << ": " << ans[i].x << "," << ans[i].y << endl;
}
//cout << getDistance(p1, p2) << endl;
}
REF: