K Closest Points to Origin

Q: Why a max-heap when we want the smallest distances?

The heap holds the current best k, and the question for each new point is "is it closer than the worst of my best?". The worst of the best is the farthest one, so it has to be on top, which is a max-heap.

Q: Does this connect to infra work?

"Closest" does not have to mean geometry. Picking the k lowest-latency regions or endpoints for a client, or the k hosts with the smallest load score, is the same bounded top-k selection with a different distance function.

mediumHeaps and top-K Must-do

You get a list of points on a flat plane, each written as [x, y], and an integer k. Return the k points that are closest to the origin (0, 0), using ordinary straight-line (Euclidean) distance.

You can return the points in any order. The input is chosen so that the set of k closest points is unambiguous.

Examples

Example 1

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

Output: [[1, -1], [-2, 2]]

Explanation: Squared distances are 25, 2 and 8. The two smallest are 2 and 8.

Example 2

Input: points = [[0, 5], [4, 4], [-1, -1], [6, 0]], k = 1

Output: [[-1, -1]]

Explanation: Squared distances are 25, 32, 2 and 36, so [-1, -1] is the closest.

Hints

Approach

Keep a bounded max-heap of the best k candidates.

For each point, compute its squared distance and push it into a max-heap keyed on that distance. In Python, push (-dist, x, y) because heapq is a min-heap.
If the heap now has more than k points, pop the top: it is the farthest of the k + 1, so it cannot be in the answer.
After all points, the heap holds exactly the k closest. Return them.

The heap never grows past k + 1, so each push and pop costs O(log k). This also works when points arrive one at a time and you cannot store them all.

ComplexityTime O(n log k)Space O(k)

Python

import heapq
 
class Solution:
    def kClosest(self, points: list[list[int]], k: int) -> list[list[int]]:
        heap = []   # max-heap by distance (stored negated), size <= k
        for x, y in points:
            dist = x * x + y * y
            heapq.heappush(heap, (-dist, x, y))
            if len(heap) > k:
                heapq.heappop(heap)          # evict the farthest of the k+1
        return [[x, y] for _, x, y in heap]

Follow-up questions

Return the points sorted from nearest to farthest.
Points arrive as an endless stream. Report the current k closest after each new point.

Frequently asked questions

The heap holds the current best k, and the question for each new point is "is it closer than the worst of my best?". The worst of the best is the farthest one, so it has to be on top, which is a max-heap.

Yes: quickselect on squared distance, the same idea as Kth Largest Element in an Array. It partitions the points so the k closest end up at the front, in O(n) on average. It needs all points in memory and has an O(n^2) worst case, so the heap is usually the first answer.

"Closest" does not have to mean geometry. Picking the k lowest-latency regions or endpoints for a client, or the k hosts with the smallest load score, is the same bounded top-k selection with a different distance function.

Problem statement