Linkedin Interview - K closest points

Find the K closest points to the origin in a 2D plane, given an array containing N points.

 

Max Heap: O(NlogN)

Selection Algorithm: O(N)

下面给出基于Max Heap的实现。

/*
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 tmp = pq.peek();
            if ((p[i].x * p[i].x + p[i].y * p[i].y) - (tmp.x * tmp.x + tmp.y * tmp.y) < 0) {
                pq.poll();
                pq.offer(p[i]);
            }
        }
    }
     
    List<Point> x = new ArrayList<>();
    while (!pq.isEmpty())
        x.add(pq.poll());
     
    return x;
}

 

Reference:

https://shepherdyuan.wordpress.com/2014/07/23/linkedin-k-closest-points/