[LintCode] K Closest Points

Problem

Given some points and a point origin in two dimensional space, find k points out of the some points which are nearest to origin.
Return these points sorted by distance, if they are same with distance, sorted by x-axis, otherwise sorted by y-axis.

Example

Given points = [[4,6],[4,7],[4,4],[2,5],[1,1]], origin = [0, 0], k = 3
return [[1,1],[2,5],[4,4]]

Solution

/**
 * 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
     */
    public Point[] kClosest(Point[] points, Point origin, int k) {
        Comparator<Point> pointComparator = new Comparator<Point>() {
            public int compare(Point A, Point B) {
                if (distance(B, origin) == distance(A, origin)) {
                    if (A.x == B.x) {
                        return A.y - B.y;
                    } else {
                        return A.x - B.x;
                    }
                } else {
                    //maintain a min heap
                    return distance(A, origin) > distance(B, origin) ? 1 : -1;
                }
            }
        };
        
        PriorityQueue<Point> Q = new PriorityQueue<Point>(k, pointComparator);
        
        for (Point point: points) {
            Q.add(point);
        }
        Point[] res = new Point[k];
        for (int i = 0; i < k; i++) {
            res[i] = Q.poll();
        }
        return res;
    }
    
    public double distance(Point A, Point B) {
        int l = Math.abs(A.x - B.x);
        int w = Math.abs(A.y - B.y);
        return Math.sqrt(Math.pow(l, 2) + Math.pow(w, 2));
    }
}