[LeetCode] 973. K Closest Points to Origin

 Problem : https://leetcode.com/problems/k-closest-points-to-origin/

Use max priority queue to find the K closest points.

Time complexity = O( N * Log K ),  N = length of the input array. 

Space complexity = O ( K ),  the priority queue will contain at most K elements.

class Solution {
    public int[][] kClosest(int[][] points, int k) {
        // create a max priority queue for saving the K closest points
        PriorityQueue<int[]> pq = new PriorityQueue<int[]>((a, b) -> distanceOf(b) - distanceOf(a));
        
        for (int i = 0; i < points.length; i++) {
            if (pq.size() < k) {
                // put the point into the priority queue if the queue size is less than K
                pq.offer(points[i]);
            } else if (distanceOf(pq.peek()) > distanceOf(points[i])) {
                // otherwise, add the new point into priority queue if it is closer than the top point.
                // because it is max priority queue, the top point is the furthest point so far.
                // this condition check is needed to avoid blindly adding new point inito priority queue.
                pq.poll();
                pq.offer(points[i]);
            }
        }
    
        return pq.toArray(new int[0][]);
    }
    
    int distanceOf(int[] p) {
        return p[0] * p[0] + p[1] * p[1];
    }
}

Quick Selection solutioin.

Time complexity = O(N)

class Solution {
    public int[][] kClosest(int[][] points, int k) {
        return quickSelect(points, 0, points.length-1, k);
    }
    
    int[][] quickSelect(int[][] points, int left, int right, int k) {       
        while (left < right) {
            int p = partition(points, left, right);
            if (p == k-1) {
                break;
            }
            
            if (p < k-1) {
                // continue to sort the rang [p+1, right]
                left = p + 1;
            } else {
                // continue to sort the range [left, p-1]
                right = p - 1;
            }
        }
        
        return Arrays.copyOf(points, k); 
    }
    
    int partition(int[][] points, int left, int right) {
        // use points[left] as the pivot value
        int pivotDistance = distanceOf(points[left]);
        
        int i = left, j = right + 1; // left and right scan indices
        
        while (true) {
            // scan left, stop when points[i] > pivotDistance or i reaches the array end
            while (distanceOf(points[++i]) < pivotDistance) if (i == right) break;
            // scan right, stop when points[j] < pivotDistance or j reaches the array start
            while (distanceOf(points[--j]) > pivotDistance) if (j == left) break;
            
            // stop scan
            if (i >= j) break;
            
            // exchange i and j
            int[] tmp = points[j];
            points[j] = points[i];
            points[i] = tmp;
        }
        
        // points[j] is the last element on right smaller than points[left]
        // swap 'left' and 'j'
        int[] tmp = points[left];
        points[left] = points[j];
        points[j] = tmp;
        
        return j;
    }
    
    int distanceOf(int[] p) {
        return p[0] * p[0] + p[1] * p[1];
    }
}

Edited on 12/26/2021. Add the quick select solution.