K Closest Points to Origin - Leetcode Solution

Problem Link

💡 Step-by-Step Thought Process

Understand the problem: Find the k points closest to the origin (0,0) from a list of points, using Euclidean distance.
Define a function to compute the squared Euclidean distance (x² + y²) for a point (x, y).
Initialize an empty max heap to store up to k points, storing tuples of (-distance, x, y).
For each point, compute its distance; if the heap has fewer than k points, push the point; otherwise, push and pop to keep the k closest.
Extract the x, y coordinates from the heap and return them as the result.

Code Solution

class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        def dist(x, y):
            return x**2 + y**2

        heap = []
        for x, y in points:
            d = dist(x, y)
            if len(heap) < k:
                heapq.heappush(heap, (-d, x, y))
            else:
                heapq.heappushpop(heap, (-d, x, y))

        return [(x, y) for d, x, y in heap]

# Time Complexity: O(n log k)
# Space Complexity: O(k)

#include <vector>
#include <queue>
#include <functional>

using namespace std;

class Solution {
public:
    vector<vector<int>> kClosest(vector<vector<int>>& points, int k) {
        auto dist = [](const vector<int>& point) {
            return point[0] * point[0] + point[1] * point[1];
        };
        
        priority_queue<vector<int>, vector<vector<int>>, function<bool(const vector<int>&, const vector<int>&)>> maxHeap(
            [&dist](const vector<int>& a, const vector<int>& b) {
                return dist(a) < dist(b);
            }
        );

        for (const auto& point : points) {
            maxHeap.push(point);
            if (maxHeap.size() > k) {
                maxHeap.pop();
            }
        }

        vector<vector<int>> result;
        while (!maxHeap.empty()) {
            result.push_back(maxHeap.top());
            maxHeap.pop();
        }
        return result;
    }
};

import java.util.PriorityQueue;
import java.util.Arrays;

class Solution {
    public int[][] kClosest(int[][] points, int k) {
        // Custom comparator to sort based on distance
        PriorityQueue<int[]> maxHeap = new PriorityQueue<>((a, b) -> {
            int distA = a[0] * a[0] + a[1] * a[1];
            int distB = b[0] * b[0] + b[1] * b[1];
            return Integer.compare(distB, distA); // Max-heap
        });

        for (int[] point : points) {
            maxHeap.add(point);
            if (maxHeap.size() > k) {
                maxHeap.poll(); // Remove the farthest point
            }
        }

        // Convert the heap to an array
        int[][] result = new int[k][2];
        for (int i = 0; i < k; i++) {
            result[i] = maxHeap.poll();
        }

        return result;
    }
}

/*
max heap approach:
have a max heap of size k, so we would do N insertions that take log(k)
for this case we would need to implement heapify up (insert) and heapify down (remove)

runtime: O(N log(k))
space: O(k)
*/
var kClosest = function(points, k) {
    let heap = []
    
    // now we need to try to add all points to the heap
    for(let i=0; i<points.length; i++) {
        if(heap.length >= k && distance(points[i]) > distance(heap[0])) { // it's bigger than the max, we can just skip it
            continue
        }
        add(heap, points[i], distance)
        if(heap.length > k) {
            remove(heap, distance)
        }
    }
    
    return heap
    
    // add at end, heapify up
    function add(heap, node, weightFunction) {
        heap.push(node)
        heapifyUp(heap, heap.length - 1, weightFunction)
    }
    
    // compare with parent and swap if needed, repeat
    function heapifyUp(heap, idx, weightFunction) {
        if(idx === 0) return
        let parent = Math.floor((idx-1) / 2)
        if(weightFunction(heap[idx]) > weightFunction(heap[parent])) {
            [heap[idx], heap[parent]] = [heap[parent], heap[idx]]
            heapifyUp(heap, parent, weightFunction)
        }
    }
    
    // read 0, replace 0 with last position, heapifyDown
    function remove(heap, weightFunction) {
        let val = heap[0]
        heap[0] = heap.pop()
        heapifyDown(heap, 0, weightFunction)
        return val
    }
    
    // compare with children, swap with biggest, repeat
    function heapifyDown(heap, idx, weightFunction) {
        let left = 2*idx + 1
        let right = 2*idx + 2
        let biggest = left
        
        if(left >= heap.length) return
        
        if(right < heap.length && weightFunction(heap[left]) < weightFunction(heap[right])) {
            biggest = right
        }
        
        if (weightFunction(heap[idx]) < weightFunction(heap[biggest])) {
            [heap[idx], heap[biggest]] = [heap[biggest], heap[idx]]
            heapifyDown(heap, biggest, weightFunction)
        }
    }
    
    function distance(point) {
        return point[0] * point[0] + point[1] * point[1]
    }
}

Detailed Explanation

Understanding the Problem: K Closest Points to Origin

The “K Closest Points to Origin” problem asks us to find the k points nearest to the origin (0,0) from a given list of 2D points. The distance is measured using the Euclidean formula, but since we only care about comparing distances (not actual values), we can safely use the squared distance formula to avoid unnecessary square root calculations.

This type of problem is common in computational geometry, spatial indexing, and real-world systems like nearest-neighbor searches in maps or recommendation engines.

Brute-Force Strategy with Max Heap

To solve the problem efficiently, we use a max heap to keep track of the k closest points. As we iterate through each point, we calculate its squared distance from the origin. We maintain a heap of size at most k containing the closest points seen so far.

To simulate a max heap using Python’s built-in heapq module (which is a min heap), we store the negative of the squared distance. If the heap grows beyond size k, we remove the point with the largest distance. At the end, the heap will contain exactly the k points closest to the origin.

Why Squared Distance?

The standard Euclidean distance formula is √(x² + y²). However, since square roots are monotonic, the relative order of distances is preserved even without computing the square root. That’s why we use the x² + y² value instead — it's simpler and faster.

Time and Space Complexity

The time complexity of this heap-based approach is O(n log k), where n is the number of points and k is the number of closest points we want. Each insertion into the heap takes O(log k), and we do this for all points.

The space complexity is O(k) for the heap that stores the closest points.

Alternative Approach: Sorting

An alternate solution is to sort the list of points by their squared distances and return the first k. This takes O(n log n) time due to the sorting step, and O(n) space if not sorting in place. While this approach is simpler, the heap-based solution is more efficient when k is much smaller than n.

Conclusion

The “K Closest Points to Origin” problem demonstrates the power of heaps in reducing the overhead of sorting when only a subset of values is needed. Using a max heap in combination with squared Euclidean distance gives us a clean, efficient solution suitable for large-scale datasets and performance-sensitive applications.

Get Personalized Support at AlgoMap Bootcamp 💡