DSA K Closest Points to the Origin

The problem

K Closest Points asks for the k points nearest to the origin. The origin is the point (0, 0). A point is a pair of coordinates (x, y), and coordinates may be negative.

Distance from the origin decides the order. This article orders by squared distance, the value x*x + y*y. Equal distances break ties by the point itself, so every run prints the same order.

def dist2(point):
    x, y = point
    return x * x + y * y

points = [(1, 3), (-2, 2), (5, 8), (0, 1), (2, 2), (3, 3)]
cases = [
    (1, [(1, (0, 1))]),
    (3, [(1, (0, 1)), (8, (-2, 2)), (8, (2, 2))]),
    (4, [(1, (0, 1)), (8, (-2, 2)), (8, (2, 2)), (10, (1, 3))]),
]
for k, expected in cases:
    print("k =", k, "->", expected)
print("squared distances:", [(p, dist2(p)) for p in points])

Each answer pair shows a squared distance and its point. The point (5, 8) scores 89 and never comes close. The two points scoring 8 print in point order, (-2, 2) before (2, 2).

Intuition: grid distance picks the wrong neighbour

The first idea measures closeness by grid distance, the value |x| + |y| along the streets of a grid. Grid distance is easy to add up. It is also the wrong measure for straight-line closeness, because a diagonal step shortens the true distance.

The trap shows on two points. The point (3, 0) wins on grid distance. The point (2, 2) is truly nearer, because its squared distance is smaller.

def dist2(point):
    x, y = point
    return x * x + y * y

a = (3, 0)
b = (2, 2)
print("grid distance:", abs(a[0]) + abs(a[1]), "vs", abs(b[0]) + abs(b[1]), "- grid picks", a)
print("squared distance:", dist2(a), "vs", dist2(b), "- truly nearer:", b)

Squared distance also removes the square roots. It orders points exactly as true distance does, because a larger true distance always has a larger square. The rest of this article never takes a square root.

Brute force: score every point, sort every point

The brute force computes the squared distance of every point. It then sorts all the points by distance and takes the first k. The sort below is a selection sort, and a counter records its comparisons.

Selection sort repeatedly moves the smallest remaining entry to the front. Sorting everything is honest work. It also orders the far points that the answer will never include.

def k_closest_brute(points, k):
    scored = []
    for point in points:
        x, y = point
        scored.append((x * x + y * y, point))
    comparisons = 0
    for i in range(len(scored)):
        best = i
        for j in range(i + 1, len(scored)):
            comparisons += 1
            if scored[j] < scored[best]:
                best = j
        scored[i], scored[best] = scored[best], scored[i]
    return scored[:k], comparisons

points = [(1, 3), (-2, 2), (5, 8), (0, 1), (2, 2), (3, 3)]
result, comparisons = k_closest_brute(points, 3)
print("k = 3 answer:", result)
print("comparisons to sort all points:", comparisons)

Fifteen comparisons sorted six points so three could be read. The comparison count grows with the square of the point count. The heap design below keeps only k points and never sorts the rest.

The optimal solution: a bounded heap holding the k nearest so far

The method in one sentence: keep a heap of at most k points keyed by negated squared distance, so the heap top is always the farthest point currently kept. Bounded means the heap size never grows past k.

A nearer arrival replaces the top. The entries also carry the negated point, so equal distances evict in a fixed order and the printed result never varies between runs.

import heapq

def k_closest(points, k):
    heap = []
    pushes = 0
    for point in points:
        x, y = point
        entry = (-(x * x + y * y), (-x, -y), point)
        if len(heap) < k:
            heapq.heappush(heap, entry)
            pushes += 1
        elif entry > heap[0]:
            heapq.heapreplace(heap, entry)
            pushes += 1
    result = sorted((-d, p) for d, _, p in heap)
    return result, pushes

points = [(1, 3), (-2, 2), (5, 8), (0, 1), (2, 2), (3, 3)]
cases = [
    (1, [(1, (0, 1))]),
    (3, [(1, (0, 1)), (8, (-2, 2)), (8, (2, 2))]),
    (4, [(1, (0, 1)), (8, (-2, 2)), (8, (2, 2)), (10, (1, 3))]),
]
for k, expected in cases:
    result, pushes = k_closest(points, k)
    print("k =", k, "->", result, "match:", result == expected, "pushes:", pushes)

All three cases pass. The far point (5, 8) entered while the heap was filling and left again when (0, 1) arrived. Points farther than the current top, like (3, 3) at the end, never enter at all.

Walkthrough: k = 3 over the six points

The walk visits the points in input order. The state is the heap, printed as (distance, point) pairs in rising distance order for reading. The farthest kept point is the last pair in each row.

import heapq

points = [(1, 3), (-2, 2), (5, 8), (0, 1), (2, 2), (3, 3)]
k = 3
heap = []
for point in points:
    x, y = point
    entry = (-(x * x + y * y), (-x, -y), point)
    if len(heap) < k:
        heapq.heappush(heap, entry)
        action = "kept, heap still filling"
    elif entry > heap[0]:
        heapq.heapreplace(heap, entry)
        action = "kept, farthest evicted"
    else:
        action = "turned away"
    kept = sorted((-d, p) for d, _, p in heap)
    print("point", point, "dist2", x * x + y * y, "-", action, "- kept:", kept)
print("answer:", sorted((-d, p) for d, _, p in heap))

The point (0, 1) evicts (5, 8), the farthest at that moment. The point (2, 2) then evicts (1, 3). The last point, (3, 3), scores 18 against a farthest kept score of 8, so it is turned away.

Complexity

The brute force scores every point and sorts them all, and its selection sort makes n times n minus 1, divided by 2, comparisons, so its time is O(n^2); a library sort would still pay O(n log n) to order points the answer discards. The heap walk makes one pass, and each push works inside a heap of size k, so its time is O(n log k).

The brute force stores a scored entry for every point, so its extra space is O(n). The bounded heap stores at most k entries, so its extra space is O(k).

for n in (6, 48, 384):
    comparisons = n * (n - 1) // 2
    print("points", n, "- sort comparisons:", comparisons, " heap size capped at: 3")

Counts beat clocks here. At 384 points the full sort pays 73,536 comparisons. The heap still holds 3 points, and no push climbs more than 2 levels inside it.

Edge cases and pitfalls

A k of zero answers an empty list. A k beyond the point count answers every point, in order, rather than failing. An empty point list answers an empty list for any k.

Duplicate points are separate points and both can be kept. A point at the origin scores zero and beats everything. Negative coordinates square away, so (-2, 2) and (2, 2) tie and print in point order.

import heapq

def k_closest(points, k):
    if k <= 0:
        return []
    heap = []
    for point in points:
        x, y = point
        entry = (-(x * x + y * y), (-x, -y), point)
        if len(heap) < k:
            heapq.heappush(heap, entry)
        elif entry > heap[0]:
            heapq.heapreplace(heap, entry)
    return sorted((-d, p) for d, _, p in heap)

print("k = 0:", k_closest([(1, 1)], 0))
print("k beyond count:", k_closest([(2, 2), (1, 1)], 5))
print("no points:", k_closest([], 3))
print("origin wins:", k_closest([(3, 3), (0, 0), (1, 0)], 2))
print("duplicates kept:", k_closest([(1, 1), (1, 1), (4, 4)], 2))
print("mirror tie:", k_closest([(2, 2), (-2, 2)], 2))

The mirror tie prints (-2, 2) before (2, 2). The negated point inside each entry fixed that order when the points entered. No run of this code can print them the other way round.