K Closest Points to Origin

Problem Statement

LeetCode-973: Given an array of points where $points[i] = [x_i,\enspace y_i]$ represents a point on the X-Y plane and an integer k, return the k closest points to the origin (0, 0).

The distance between two points on the X-Y plane is the Euclidean distance i.e., $\sqrt{(x_1 - x_2)^2 + (y_1 - y_2)^2}$. You may return the answer in any order. The answer is guaranteed to be unique (except for the order that it is in).

Approach

Problems that require finding the k smallest or largest elements are generally a good candidate for using a heap. In this problem, we need to find the k closest points to the origin. We can calculate the distance of each point from the origin and then store these distances in a min heap.

Once done, the top k elements of the heap will contain the k closest points to the origin. We can then return these points.

Implementation

Here is a Java implementation of the above approach:

public int[][] kClosest(int[][] points, int k) {
    // min heap to store the points based on their distance from the origin
    // as the co-od for origin is (0, 0), the x2 and y2 
    // are simply ignored in the distance calculation
    Queue<int[]> minHeap = new PriorityQueue<>((one, two) -> {
        double d1 = Math.sqrt((one[0]*one[0]) + (one[1]*one[1]));
        double d2 = Math.sqrt((two[0]*two[0]) + (two[1]*two[1]));
        return Double.compare(d1, d2);
    });
 
    for(int[] point : points){
        minHeap.add(point);
    }
    
    List<int[]> result = new ArrayList<>();
    // get the top k closest points
    while(k > 0){
        result.add(minHeap.remove());
        k--;
    }

    return result.toArray(new int[0][0]);
}

Complexity Analysis

The time complexity for this approach is $O(n\cdot \log n)$ where n is the number of points. This is because we are iterating over all the points and adding them to the heap which takes $O(\log n)$ time. The space complexity is $O(n)$ as we are storing all the points in the heap.

Follow-up

Let us see if we can improve the time and space complexity of the above approach. If we revisit the problem, we can see that we are only interested in the k closest points to the origin, but we are still storing all the points in the heap.

Can we do better? Let’s see if we can only store the k closest points in the minHeap and discard the rest.

1. Add the current point to the heap.
2. If the size of the heap exceeds k, remove the top of the heap so that the heap only contains the k closest points.

But this approach will not work as the top of the minHeap contains the closest point to the origin. So, we need to use a maxHeap instead of a minHeap. This way, the top of the maxHeap will contain the farthest point from the origin.

Here is a Java implementation of the above approach:

public int[][] kClosest(int[][] points, int k) {
    Queue<int[]> maxHeap = new PriorityQueue<>((one, two) -> {
        double d1 = Math.sqrt((one[0]*one[0]) + (one[1]*one[1]));
        double d2 = Math.sqrt((two[0]*two[0]) + (two[1]*two[1]));
        // compare in reverse order
        return Double.compare(d2, d1); 
    });

    for(int[] point : points){
        minHeap.add(point);

        // if the size of the heap exceeds k, 
        // remove the top of the heap
        if(minHeap.size() > k){
            minHeap.remove();
        }
    }
    return minHeap.toArray(new int[0][0]);
}

The time complexity for this approach is $O(n\cdot \log k)$ where n is the number of points and k is the number of closest points we are interested in. The space complexity is $O(k)$ as we are storing only the k closest points in the heap.