Module 09 · Queues, Heaps, and Streaming Frontiers

089. K Closest Points to Origin

Interview · Download this solution · All solutions and shared helpers

Problem

Return up to k two-dimensional integer points closest to the origin, ordered by squared distance and then original input index. Duplicated points count as distinct occurrences. Return coordinate tuples; k<=0 yields []. Do not use square roots or mutate the input.

Examples

Example 1

closest_points([(1,3),(-2,2)], 1)
=> [(-2, 2)]

Squared distance eight is smaller than ten.

Example 2

closest_points([(1,0),(0,1),(-1,0)], 2)
=> [(1, 0), (0, 1)]

Equal distances retain the earliest input indices.

Example 3

closest_points([], 3)
=> []

There are no points to retain.

Reveal the worked solution and reasoning

Why this mechanism works

Distance ordering is unchanged by squaring, avoiding floating-point comparisons. Retain the k strongest candidates in a max-oriented heap represented by negative distance and negative index. The worst retained distance, with the latest tied index, is evicted first.

State and boundaries

heap stores (-squared_distance,-input_index) for retained points. The root is the least desirable retained point.

Protocol

Compute x*x+y*y.
Push the negated ranking key.
Evict one root when capacity is exceeded.
Sort retained negative keys descending.
Recover coordinate tuples by saved index.

Visualization

(1,3): squared distance 10
(-2,2): squared distance 8
retain the latter for k=1

Python implementation

import heapq

def closest_points(points, k):
    if k <= 0:
        return []
    heap = []
    for i, (x, y) in enumerate(points):
        heapq.heappush(heap, (-(x*x + y*y), -i))
        if len(heap) > k:
            heapq.heappop(heap)
    return [tuple(points[-index]) for distance, index in sorted(heap, reverse=True)]

Worked execution

Recorded state changes after statements; L numbers refer to the solution file. Long traces are explicitly condensed.

closest_points([(1,3),(-2,2)], 1)

L4 closest_points: if k <= 0:
points=[(1, 3), (-2, 2)]; k=1

L6 closest_points: heap = []
points=[(1, 3), (-2, 2)]; k=1; heap=[]

L7 closest_points: for i, (x, y) in enumerate(points):
points=[(1, 3), (-2, 2)]; k=1; heap=[]; i=0; x=1; y=3

L8 closest_points: heapq.heappush(heap, (-(x*x + y*y), -i))
points=[(1, 3), (-2, 2)]; k=1; heap=[(-10, 0)]; i=0; x=1; y=3

L7 closest_points: for i, (x, y) in enumerate(points):
points=[(1, 3), (-2, 2)]; k=1; heap=[(-10, 0)]; i=1; x=-2; y=2

L8 closest_points: heapq.heappush(heap, (-(x*x + y*y), -i))
points=[(1, 3), (-2, 2)]; k=1; heap=[(-10, 0), (-8, -1)]; i=1; x=-2; y=2

L10 closest_points: heapq.heappop(heap)
points=[(1, 3), (-2, 2)]; k=1; heap=[(-8, -1)]; i=1; x=-2; y=2

L11 closest_points: return [tuple(points[-index]) for distance, index in sorted(heap, reverse=True)]
points=[(1, 3), (-2, 2)]; k=1; heap=[(-8, -1)]; i=1; x=-2; y=2; distance=-8; index=-1

L11 closest_points: return [tuple(points[-index]) for distance, index in sorted(heap, reverse=True)]
points=[(1, 3), (-2, 2)]; k=1; heap=[(-8, -1)]; i=1; x=-2; y=2

Result: [(-2, 2)]. Squared distance eight is smaller than ten.

closest_points([(1,0),(0,1),(-1,0)], 2)

L4 closest_points: if k <= 0:
points=[(1, 0), (0, 1), (-1, 0)]; k=2

L6 closest_points: heap = []
points=[(1, 0), (0, 1), (-1, 0)]; k=2; heap=[]

L7 closest_points: for i, (x, y) in enumerate(points):
points=[(1, 0), (0, 1), (-1, 0)]; k=2; heap=[]; i=0; x=1; y=0

L8 closest_points: heapq.heappush(heap, (-(x*x + y*y), -i))
points=[(1, 0), (0, 1), (-1, 0)]; k=2; heap=[(-1, 0)]; i=0; x=1; y=0

L7 closest_points: for i, (x, y) in enumerate(points):
points=[(1, 0), (0, 1), (-1, 0)]; k=2; heap=[(-1, 0)]; i=1; x=0; y=1

L8 closest_points: heapq.heappush(heap, (-(x*x + y*y), -i))
points=[(1, 0), (0, 1), (-1, 0)]; k=2; heap=[(-1, -1), (-1, 0)]; i=1; x=0; y=1

L7 closest_points: for i, (x, y) in enumerate(points):
points=[(1, 0), (0, 1), (-1, 0)]; k=2; heap=[(-1, -1), (-1, 0)]; i=2; x=-1; y=0

L8 closest_points: heapq.heappush(heap, (-(x*x + y*y), -i))
points=[(1, 0), (0, 1), (-1, 0)]; k=2; heap=[(-1, -2), (-1, 0), (-1, -1)]; i=2; x=-1; y=0

L10 closest_points: heapq.heappop(heap)
points=[(1, 0), (0, 1), (-1, 0)]; k=2; heap=[(-1, -1), (-1, 0)]; i=2; x=-1; y=0

L11 closest_points: return [tuple(points[-index]) for distance, index in sorted(heap, reverse=True)]
points=[(1, 0), (0, 1), (-1, 0)]; k=2; heap=[(-1, -1), (-1, 0)]; i=2; x=-1; y=0; distance=-1; index=0

L11 closest_points: return [tuple(points[-index]) for distance, index in sorted(heap, reverse=True)]
points=[(1, 0), (0, 1), (-1, 0)]; k=2; heap=[(-1, -1), (-1, 0)]; i=2; x=-1; y=0; distance=-1; index=-1

L11 closest_points: return [tuple(points[-index]) for distance, index in sorted(heap, reverse=True)]
points=[(1, 0), (0, 1), (-1, 0)]; k=2; heap=[(-1, -1), (-1, 0)]; i=2; x=-1; y=0

Result: [(1, 0), (0, 1)]. Equal distances retain the earliest input indices.

Why the solution is correct

Squared distances preserve nonnegative Euclidean-distance order. Negated keys expose the worst winner at the min-heap root, including the later-index tie rule. Eviction retains exactly the best k keys, and final sorting restores the required closest-first order.

Time and space

O(n log(k+1)+k log k) time and O(k) working space under unit-cost integer arithmetic.

Common mistakes

Square roots add unnecessary rounding risk. Comparing point coordinates to resolve ties does not implement the stated input-index policy.

Reasoning card

State: heap stores (-squared_distance,-input_index) for retained points. The root is the least desirable retained point.

Invariant: Retained points are the best k among those processed.

Decision to rehearse: Use a bounded max-oriented heap or a full min-heap deliberately.

Finish: Recover coordinate tuples by saved index

Transfer exercise

Use quickselect to partition the nearest k points, and distinguish selection cost from the cost of ordered output.