Kth closest points to origin

Problem Link

K Closest Points to Origin – LeetCode

Problem: K Closest Points to Origin

Problem Statement

Given an array of points where points[i] = [xi, yi] 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 (x1, y1) and (x2, y2) is given by the Euclidean distance formula:
√((x1 - x2)² + (y1 - y2)²)

You may return the answer in any order.

Examples

Example 1
Input: points = [[1,3],[-2,2]], k = 1
Output: [[-2,2]]
Explanation: The distance of (1,3) is √10. The distance of (-2,2) is √8. Since √8 < √10, we return [-2,2].

Example 2
Input: points = [[3,3],[5,-1],[-2,4]], k = 2
Output: [[3,3],[-2,4]]
Explanation: The distances are: √18, √26, √20. The two smallest are √18 and √20.

Constraints

• 1 <= k <= points.length <= 10⁴

• -10⁴ < xi, yi < 10⁴

Intuition

We need to extract the k closest points to the origin based on the Euclidean distance. Instead of calculating the actual square root (which is computationally expensive), we can compare using the squared distances. To efficiently track the k closest points while traversing all points, a max-heap of size k is ideal.

Approach

1. Use a max-heap (priority queue in descending order of squared distance to origin).

2. For each point:

   • Calculate its squared distance: x² + y².

   • Add it to the heap.

   • If the heap size exceeds k, remove the farthest point (root of the max-heap).

3. After processing all points, extract the elements from the heap to form the result array.

Code (Java)

import java.util.*;

class Solution {
    public int[][] kClosest(int[][] points, int k) {
        PriorityQueue<int[]> pq = new PriorityQueue<>(
            (a, b) -> (b[0]*b[0] + b[1]*b[1]) - (a[0]*a[0] + a[1]*a[1])
        );
        
        for (int[] point : points) {
            pq.add(point);
            if (pq.size() > k) {
                pq.poll();
            }
        }
        
        int[][] ans = new int[k][2];
        for (int i = 0; i < k; i++) {
            int[] polled = pq.poll();
            ans[i][0] = polled[0];
            ans[i][1] = polled[1];
        }
        
        return ans;
    }
}

Time and Space Complexity

• Time Complexity: O(n log k)
  Each insertion and removal in the heap of size k takes O(log k), and we process n points.

• Space Complexity: O(k)
  The heap stores up to k points at any time. The result array also takes O(k) space.

Dry Run

Input: points = [[3,3],[5,-1],[-2,4]], k = 2

• Initialize max-heap based on squared distance

1. Add [3,3] → distance² = 18 → heap = [[3,3]]

2. Add [5,-1] → distance² = 26 → heap = [[5,-1],[3,3]]

3. Add [-2,4] → distance² = 20

   • Heap size > 2 → remove farthest (26) → heap = [[3,3],[-2,4]]

Extract and return points from heap → Output: [[3,3],[-2,4]]

Conclusion

This is a classic example of using a max-heap to maintain a dynamic list of the k closest elements. Comparing squared distances eliminates the need for square roots and speeds up computations. The approach guarantees efficient performance for large datasets.