描述

给定一些points和一个origin，从points中找到k个离origin最近的点。按照距离由小到大返回。如果两个点有相同距离，则按照x值来序；若x值也相同，就再按照y值排序。

样例

给出 points = [[4,6],[4,7],[4,4],[2,5],[1,1]], origin = [0, 0], k = 3
返回 [[1,1],[2,5],[4,4]]

代码

/**
 * Definition for a point.
 * class Point {
 *     int x;
 *     int y;
 *     Point() { x = 0; y = 0; }
 *     Point(int a, int b) { x = a; y = b; }
 * }
 */
public class Solution {
    /**
     * @param points a list of points
     * @param origin a point
     * @param k an integer
     * @return the k closest points
     */
    private Point global_origin = null;
    public Point[] kClosest(Point[] points, Point origin, int k) {
        global_origin = origin;
        // 由大到小顺序排列，最大堆来实现PriorityQueue
        PriorityQueue<Point> pq = new PriorityQueue<Point> (k, new Comparator<Point> () {
            @Override
            public int compare(Point a, Point b) {
                int diff = getDistance(b, global_origin) - getDistance(a, global_origin);
                if (diff == 0) {
                    diff = b.x - a.x;
                }
                if (diff == 0) {
                    diff = b.y - a.y;
                }
                return diff;
            }
        });

        // 遍历整个数组，把前k小的数加入PriorityQueue
        for (int i = 0; i < points.length; i++) {
            pq.offer(points[i]);
            if (pq.size() > k) {
                pq.poll();
            }
        }
        
        k = pq.size();
        Point[] ret = new Point[k];  
        // 最大堆先抛出来的是最大的，
        // 所以要让ret[k - 1]等于最大堆抛出的第一个元素,依次赋值
        while (!pq.isEmpty()) {
            ret[--k] = pq.poll();
        }
        return ret;
    }
    
    private int getDistance(Point a, Point b) {
        return (a.x - b.x) * (a.x - b.x) + (a.y - b.y) * (a.y - b.y);
    }
}