K Closest Points

Java

public class Solution {
    /**
     * @param points: a list of points
     * @param origin: a point
     * @param k: An integer
     * @return: the k closest points
     */
    public Point[] kClosest(Point[] points, Point origin, int k) {
        // write your code here
        PriorityQueue<Point> pq = new PriorityQueue<>(k, (a, b) -> {
            int diff = getDistance(b, origin) - getDistance(a, origin);
            if (diff == 0) {
                diff = b.x - a.x;
            }
            if (diff == 0) {
                diff = b.y - a.y;
            }
            return diff;
        });
        
        for (int i = 0; i < points.length; i++) {
            pq.offer(points[i]);
            if (pq.size() > k) {
                pq.poll();
            }
        }
        
        k = pq.size();
        Point[] res = new Point[k];
        while (!pq.isEmpty()) {
            k--;
            res[k] = pq.poll();
        }
        return res;
    }
    
    public int getDistance(Point a, Point b) {
        return (a.x-b.x)*(a.x-b.x) + (a.y-b.y)*(a.y-b.y);
    }
}