Java: K Closest Points to Origin

Thought Process

To efficiently find the K closest points to the origin, we need a strategy that maintains optimal candidates. We use a max heap to maintain the K smallest distances. For each point, we calculate its squared distance (avoiding square root for efficiency) and push it into the heap. When the heap size exceeds K, we remove the top element of the heap. This ensures we always keep the K closest points in the heap. 

class Solution {
    public int[][] kClosest(int[][] points, int k) {
        
        PriorityQueue<int[]> pq = new PriorityQueue<>((a,b)->Integer.compare(b[0],a[0]));
        for(int idx=0;idx<points.length;idx++){

            int distance = points[idx][0]*points[idx][0]  + points[idx][1]*points[idx][1];
            pq.offer(new int[]{distance,idx});

            if(pq.size()>k)
                pq.poll();
        }
        List<int[]> ans = new ArrayList<>();
        while(!pq.isEmpty()){
            int idx = pq.poll()[1];
            ans.add(points[idx]);
        }
        return ans.toArray(new int[ans.size()][2]);
    }
}

Code Complexity

Time Complexity: O(n log k)

We process n points, with each heap operation taking O(log k) time since we maintain at most k elements.

Space Complexity: O(k)

The heap stores at most k elements, and the output requires O(k) space.

Code copied!