K Closest Points to Origin

Problem

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 on the X-Y plane is the Euclidean distance (i.e., √(x1 - x2)² + (y1 - y2)²).

You may return the answer in any order. The answer is guaranteed to be unique (except for the order that it is in).

Example 1:

Input: points = [[1,3],[-2,2]], k = 1
Output: [[-2,2]]
Explanation:
The distance between (1, 3) and the origin is sqrt(10).
The distance between (-2, 2) and the origin is sqrt(8).
Since sqrt(8) < sqrt(10), (-2, 2) is closer to the origin.

Example 2:

Input: points = [[3,3],[5,-1],[-2,4]], k = 2
Output: [[3,3],[-2,4]]
Explanation: The answer [[-2,4],[3,3]] would also be accepted.

Constraints:

1 <= k <= points.length <= 10^4
-10^4 < xi, yi < 10^4

Solution

Approach 1: Max Heap of Size K

The key insight is to maintain a max heap of size k with the k closest points. We use max heap so we can remove the farthest point when needed.

Implementation

import heapq

class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        # Max heap: store (-distance, point)
        heap = []

        for x, y in points:
            dist = x * x + y * y  # No need for sqrt for comparison
            heapq.heappush(heap, (-dist, [x, y]))

            # Keep only k closest
            if len(heap) > k:
                heapq.heappop(heap)

        return [point for _, point in heap]

Approach 2: Sort (Simple)

class Solution:
    def kClosest(self, points: List[List[int]], k: int) -> List[List[int]]:
        points.sort(key=lambda p: p[0] * p[0] + p[1] * p[1])
        return points[:k]

Complexity Analysis

Max Heap:

Time Complexity: O(n log k), where n is the number of points. Each heap operation is O(log k).
Space Complexity: O(k) for the heap.

Sorting:

Time Complexity: O(n log n) for sorting.
Space Complexity: O(1) or O(n) depending on sorting implementation.

Key Insights

Distance Formula: We can skip the square root since we’re only comparing distances. x² + y² preserves ordering.
Max Heap Strategy: Use a max heap so the farthest of the k closest points is at the root and can be removed efficiently.
Negative Distance: We negate distances to simulate a max heap with Python’s min heap.
Trade-offs: Max heap is better for large n and small k. Sorting is simpler and better when k is close to n.
QuickSelect: For even better average performance, QuickSelect achieves O(n) average time, but heap is simpler to implement.