K Closet Points

Given some points and a point origin in two dimensional space, find k points out of the some points which are nearest to origin(It doesn't have to be (0,0)). 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]]

input是(List<Points>, k), output(List<Points>)，有个Point class，求离原点最近的K个点。

问我如果让我做测试，我会选什么样的点，答数很大会overflow的点。然后问我如何优化时间，想了半天没想出来，然后告诉我每次push点的时候可以比一下队列顶端的点，如果更远就不用push了。。

Solution 1: max heap O(n log k); O(k)

Create a max heap to store points whose elements are ordered according to a specified comparator. The comparator is to sort points in descending order of their distance to the origin.
Put points into the max heap the size of which is k. When the size reaches k, pop the top element from the heap.
Then pop all elements from heap and store them in the result.

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

    // 方法1: max heap
    public Point[] kClosest(Point[] points, Point origin, int k) {
        PriorityQueue<Point> maxHeap = new PriorityQueue<>((a, b) -> {
            int dist = Long.compare(distance(b, origin), distance(a, origin));
            if (dist == 0) {
                dist = b.x - a.x;
            }
            if (dist == 0) { //x coordinate is equal, continue to compare y coordinate
                dist = b.y - a.y;
            }
            return dist;
        });

        for (Point p : points) {
            maxHeap.offer(p);
            if (maxHeap.size() > k) {
                maxHeap.poll();
            }
        }

        k = maxHeap.size(); //avoid that k is larger than points.length
        Point[] res = new Point[k];
        while (!maxHeap.isEmpty()) {
            res[--k] = maxHeap.poll();
        }
        return res;
    }

    public long distance(Point a, Point b) {
        return (a.x - b.x) * (a.x - b.x) + (a.y - b.y) * (a.y - b.y);
    }

Solution 2: Quick Select O(n) average, O(n + klogk) time if output is sorted, O(n^2) worst case; O(1)

Given points =[[4,6],[4,7],[4,4],[2,5],[1,1]], origin =[0, 0], k =3
return[[1,1],[2,5],[4,4]]
pivot [1,1] 第一轮后[1,1]和[4,6]交换，leftIndex = 1.
pivot [4,6]

   public Point[] kClosestB(Point[] points, Point origin, int k) {
        // write your code here
        QuickSelect(points, k, origin, 0, points.length - 1);
        Point[] res = new Point[k];
        for (int i = 0; i < k; i++) {
            res[i] = points[i];
        }
        Arrays.sort(res, (b, a) -> { //(b, a) ascending order, instead of (b, a)
            int dist = Long.compare(getDistance(b, origin), getDistance(a, origin));
            if (dist == 0) {
                dist = b.x - a.x;
            }
            if (dist == 0) { //x coordinate is equal, continue to compare y coordinate
                dist = b.y - a.y;
            }
            return dist;
        });
        return res;
    }

    public Point QuickSelect(Point[] points, int k, Point target, int start, int end) {
        if (start > end) return null;
        Point pivot = points[end];
        int leftIndex = start;
        for (int i = start; i < end; i++) {
            //核心思想：距离近的点交换到前面
            if (getDistance(points[i], target) <= getDistance(pivot, target)) 
                swap(points, leftIndex++, i);
        }
        swap(points, leftIndex, end); //然后将轴放到中间，保证左小右大
        if (leftIndex == k) return points[leftIndex];
        else if (leftIndex < k) return QuickSelect(points, k, target, leftIndex + 1, end);//在右边
        else return QuickSelect(points, k, target, start, leftIndex - 1); //在左边
    }

    private int getDistance(Point p, Point target) {
        return (p.x - target.x) * (p.x - target.x) + (p.y - target.y) * (p.y - target.y);
    }

    private static void swap(Point[] points, int left, int right) {
        Point temp = points[left];
        points[left] = points[right];
        points[right] = temp;
    }

Follow Up:

1.从N个点中找到离原点第K近的点。
讨论了直接sort, quickselect, heap各种做法的idea和时间复杂度之后，用heap的做法写这个题目。
当K的值非常大，应该怎么做。一开始没有想明白，后来经小哥善意的提醒终于给出答案。然后就是改code, 让代码能够同时适用k非常小以及k非常大。http://www.1point3acres.com/bbs/thread-228392-1-1.html

    /**
     * http://www.1point3acres.com/bbs/thread-268942-1-1.html
     * 问了那道k closest points to (0,0)，楼主用了priority queue，问了时间和空间复杂度。
     * 然后小哥继续问能不能把时间复杂度继续降低，思路跟LC215一样。小哥说不用写了，继续下一题。
     * pq的复杂度是nlogk。用quick sort的helper function而不是quick sort，因为这里不需要sort所有的点，
     * 所以average可以做到n，worst case是n^2。
     *
     * priorityqueue解 improve：用map或者建一个class，避免重复计算distance
     */

Python: 感受下代码有多短

    def kClosest(self, points, origin, k):
        def distance(p0,p1):
            return ((p0.x - p1.x)**2 + (p0.y - p1.y)**2)**0.5

        return sorted(points, key=lambda p1:(distance(origin,p1),p1.x) )[:k]
        # notice: if you want the second key p1.x in descending order, you just need to add a minus sign