K Closest Points to Origin

We have a list of points on the plane.  Find the K closest points to the origin (0, 0).

(Here, the distance between two points on a plane is the Euclidean distance.)

You may return the answer in any order.  The answer is guaranteed to be unique (except for the order that it is in.)

Example 1:

Input: points = [[1,3],[-2,2]], K = 1
Output: [[-2,2]]
Explanation: 
The distance between (1, 3) and the origin is sqrt(10).
The distance between (-2, 2) and the origin is sqrt(8).
Since sqrt(8) < sqrt(10), (-2, 2) is closer to the origin.
We only want the closest K = 1 points from the origin, so the answer is just [[-2,2]].

思路2：最佳答案是 O(N)，quick select;

class Solution {
    public int[][] kClosest(int[][] points, int k) {
        findKth(points, 0, points.length - 1, k);
        int[][] res = new int[k][2];
        for(int i = 0; i < k; i++) {
            res[i][0] = points[i][0];
            res[i][1] = points[i][1];
        }
        return res;
    }
    
    private int[] findKth(int[][] points, int start, int end, int k) {
        int i = start; int j = end;
        int mid = start + (end - start) / 2;
        int pivot = getDistance(points[mid][0], points[mid][1]);
        while(i <= j) {
            while(i <= j && getDistance(points[i][0], points[i][1]) < pivot) {
                i++;
            }
            while(i <= j && getDistance(points[j][0], points[j][1]) > pivot) {
                j--;
            }
            if(i <= j) {
                int[] temp = points[i];
                points[i] = points[j];
                points[j] = temp;
                i++;
                j--;
            }
        }
        // j, j + 1, i
        if(start + k - 1 <= j) {
            return findKth(points, start, j, k);
        }
        if(start + k - 1 >= i) {
            return findKth(points, i, end, k - (i - start));
        }
        return points[j + 1];
    }
    
    private int getDistance(int x, int y) {
        return x * x + y * y;
    }
}

思路：直接用minheap就可以解决；NlogN;

class Solution {
    public class Node {
        int x;
        int y;
        int dis;
        public Node(int x, int y, int dis) {
            this.x = x;
            this.y = y;
            this.dis = dis;
        }
    }
    
    public int[][] kClosest(int[][] points, int k) {
        PriorityQueue<Node> pq = new PriorityQueue<Node>((a, b) -> (a.dis - b.dis));
        for(int[] point: points) {
            int x = point[0];
            int y = point[1];
            pq.offer(new Node(x, y, x * x + y * y));
        }
        
        int[][] res = new int[k][2];
        int index = 0;
        while(!pq.isEmpty()) {
            Node node = pq.poll();
            res[index][0] = node.x;
            res[index][1] = node.y;
            if(index == k - 1) {
                break;
            }
            index++;
        }
        return res;
    }
}