K Closest Points to Origin

Problem

https://leetcode.com/problems/k-closest-points-to-origin/

Given an array of points where points[i] = [x_i, 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₁- x₂)²+ (y₁- y₂)²).

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.
We only want the closest k = 1 points from the origin, so the answer is just [[-2,2]].

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⁴

• -10⁴<= x_i, y_i<= 10⁴

Pattern

Heap, Sorting

Approaches

Max Heap

Explanation

The \(k`th closest point is the largest element among the :math:`k\) largest elements. So if we maintain a max-heap containing only the \(k\) closest points seen so far, we only need to look at the top of the heap to determine whether a new point should be in our \(k\) closest.

Go through points and push \((x^2 + y^2, x, y)\) onto the heap until it holds \(k\) points. After that, each new point only matters if it is closer than the top point (the furthest of the current \(k\) closest points). In that case, the top point can no longer be one of the \(k\) closest, so we replace it.

Once all numbers are processed, the heap holds the \(k\) closest points.

Code

import heapq


def kClosest(points: list[list[int]], k: int) -> list[list[int]]:
    """Return the k closest points to the origin."""

    heap = []
    for [x, y] in points:
        d = x**2 + y**2

        if len(heap) < k:
            heapq.heappush_max(heap, (d, x, y))
        elif d < heap[0][0]:
            heapq.heapreplace_max(heap, (d, x, y))

    return [[x, y] for _, x, y in heap]

Test

>>> from k_closest_points_to_origin__max_heap import kClosest
>>> kClosest([[1, 3], [-2, 2]], 1)
[[-2, 2]]
>>> sorted(kClosest([[3, 3], [5, -1], [-2, 4]], 2))
[[-2, 4], [3, 3]]

Complexity

\(n\) is the number of elements in nums

Measure	Complexity	Notes
Time	\(O(n \log k)\)	one pass through the array, pushing elements onto the heap is \(O(\log k)\)
Auxiliary Space	\(O(k)\)	max heap

k_closest_points_to_origin__max_heap.kClosest(points: list[list[int]], k: int) → list[list[int]]

Return the k closest points to the origin.