C++: 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:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int k) {
        
        // Max Heap. Hence, farthest point will be on the top of the Heap.
        priority_queue<pair<int,vector<int>>> pq;
        vector<vector<int>> ans;

        for(vector<int> entity : points){

            int x = entity[0]; // x-coordinate of point
            int y = entity[1]; // y-coordinate of point

            int distance = x*x + y*y;
            pq.push({distance,entity});

            if(pq.size()>k)
                pq.pop();
        }
        while(!pq.empty()){

            ans.push_back(pq.top().second);
            pq.pop();
        }
        return ans;
    }
};

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!