K closest points

Posted 2020-08-26 北叶青藤

Find the K closest points to a target point in a 2D plane.

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.PriorityQueue;

class Point {
    public int x;
    public int y;

    public Point(int x, int y) {
        this.x = x;
        this.y = y;
    }
}

class Solution {

    public List<Point> findKClosest(Point[] points, int k, Point p) {

        // max heap
        PriorityQueue<Point> pq = new PriorityQueue<>(10, new Comparator<Point>() {
            @Override
            public int compare(Point a, Point b) {
                double d1 = distance(a, p);
                double d2 = distance(b, p);
                if (d1 == d2)
                    return 0;
                if (d1 < d2)
                    return 1;
                return -1;
            }
        });

        for (int i = 0; i < points.length; i++) {
            pq.offer(points[i]);

            if (pq.size() > k) {
                pq.poll();
            }
        }

        List<Point> x = new ArrayList<>();
        while (!pq.isEmpty())
            x.add(pq.poll());

        return x;
    }

    private double distance(Point p1, Point p2) {
        return Math.sqrt((p1.x - p2.x) * (p1.x - p2.x) + (p1.y - p2.y) * (p1.y - p2.y));
    }
}