K Closest Points

import java.util.PriorityQueue;
import java.util.Comparator;
import java.lang.Math;

public class kNearestPoint {
    public Point[] Solution(Point[] array, Point origin, int k) {
        Point[] ans = new Point[k];  
        int index = 0;

        // 大顶堆
        PriorityQueue<Point> pq = new PriorityQueue<Point> (k, new Comparator<Point> () {
            @Override
            public int compare(Point a, Point b) {
                return (int) (getDistance(b, origin) - getDistance(a, origin));
            }
        });

        // 加入元素
        for (int i = 0; i < array.length; i++) {
            if (pq.size() < k) {
                pq.offer(array[i]);
            } else if (array[i] < pq.peek()) {
                pq.poll();
                pq.offer(array[i]);
            }
        }

        // 存入数组
        while (!pq.isEmpty())
            ans[index++] = pq.poll();
        return ans;
    ｝

    private double getDistance(Point a, Point b) {
        return (a.x - b.x) * (a.x - b.x) + (a.y - b.y) * (a.y - b.y);
    }
}