973. 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]].

Example 2:

Input: points = [[3,3],[5,-1],[-2,4]], K = 2
Output: [[3,3],[-2,4]]
(The answer [[-2,4],[3,3]] would also be accepted.)

Note:

1. 1 <= K <= points.length <= 10000
2. -10000 < points[i][0] < 10000
3. -10000 < points[i][1] < 10000

M1: min heap

time = O(n + klogn), space = O(n)

class Solution {
    public int[][] kClosest(int[][] points, int K) {
        int[][] res = new int[K][2];
        if(points == null || points.length == 0 || K == 0) {
            return res;
        }
        PriorityQueue<int[]> minHeap = new PriorityQueue<>(new Comparator<int[]>() {
            @Override
            public int compare(int[] p1, int[] p2) {
                double d1 = getDistance(p1, new int[] {0, 0});
                double d2 = getDistance(p2, new int[] {0, 0});
                if(d1 == d2) {
                    return 0;
                }
                return d1 < d2 ? -1 : 1;
            }
        });
        
        for(int[] point : points) {
            minHeap.offer(point);
        }
        for(int i = 0; i < K; i++) {
            res[i] = minHeap.poll();
        }
        
        return res;
    }
    
    public double getDistance(int[] p, int[] o) {
        int x = p[0] - o[0];
        int y = p[1] - o[1];
        return Math.sqrt(x * x + y * y);
    }
}

M2: max heap

time = O(k + (n-k)logk), space = O(k)

class Solution {
    public int[][] kClosest(int[][] points, int K) {
        int[][] res = new int[K][2];
        if(points == null || points.length == 0 || K == 0) {
            return res;
        }
        PriorityQueue<int[]> maxHeap = new PriorityQueue<>(new Comparator<int[]>() {
            @Override
            public int compare(int[] p1, int[] p2) {
                double d1 = getDistance(p1, new int[] {0, 0});
                double d2 = getDistance(p2, new int[] {0, 0});
                if(d1 == d2) {
                    return 0;
                }
                return d1 < d2 ? 1 : -1;
            }
        });
        
        for(int[] point : points) {
            if(maxHeap.size() < K) {
                maxHeap.offer(point);
            } else if(getDistance(point, new int[] {0, 0}) < getDistance(maxHeap.peek(), new int[] {0, 0})) {
                maxHeap.poll();
                maxHeap.offer(point);
            }
        }
        
        for(int i = K - 1; i >= 0; i--) {
            res[i] = maxHeap.poll();
        }
        
        return res;
    }
    
    public double getDistance(int[] p, int[] o) {
        int x = p[0] - o[0];
        int y = p[1] - o[1];
        return Math.sqrt(x * x + y * y);
    }
}

M3: quick select

time = O(n), space = O(logn)

class Solution {
    int[] origin = {0, 0};
    
    public int[][] kClosest(int[][] points, int K) {
        int[][] res = new int[K][2];
        if(points == null || points.length == 0 || K == 0) {
            return res;
        }
        quickSelect(points, 0, points.length - 1, K - 1);
        for(int i = 0; i < K; i++) {
            res[i] = points[i];
        }
        return res;
    }
    
    public void quickSelect(int[][] arr, int left, int right, int target) {
        int mid = partition(arr, left, right);
        if(mid == target) {
            return;
        } else if(mid < target) {
            quickSelect(arr, mid + 1, right, target);
        } else {
            quickSelect(arr, left, mid - 1, target);
        }
    }
    
    public int partition(int[][] arr, int left, int right) {
        int[] pivot = arr[right];
        int start = left, end = right - 1;
        while(start <= end) {
            if(getDistance(arr[start], origin) < getDistance(pivot, origin)) {
                start++;
            } else if(getDistance(arr[end], origin) >= getDistance(pivot, origin)) {
                end--;
            } else {
                swap(arr, start++, end--);
            }
        }
        swap(arr, start, right);
        return start;
    }
    
    public void swap(int[][] arr, int i, int j) {
        int[] tmp = arr[i];
        arr[i] = arr[j];
        arr[j] = tmp;
    }
    
    public double getDistance(int[] p, int[] o) {
        int x = p[0] - o[0];
        int y = p[1] - o[1];
        return Math.sqrt(x * x + y * y);
    }
}