K Closest Points to Origin Without Square Roots

"K Closest Points to Origin" shows up again and again in Meta, Amazon, Asana phone screens. It is a Medium problem on paper, but the real test is whether you recognize the Top-K and Heaps pattern quickly and code it cleanly. Part 5 of 11 in the Top-K and Heaps arc.

The Problem

Given points on a plane and an integer k, return the k points closest to the origin (any order). Two embedded lessons: skip the square root (distance comparisons survive monotone transforms), and flip the heap polarity for k-smallest problems.

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

Recognizing the Top-K and Heaps Pattern

When a problem asks for the k largest, smallest, closest, or most frequent — or for repeated access to an extreme while data changes — a heap gives O(log n) insertion and O(1) access to the extreme. The signature trick: keep a bounded heap of size k with the opposite polarity (a min-heap to track the k largest), evicting the root on overflow. Python's heapq is a min-heap; negate values for max behavior.

K smallest by a computed score — the mirror of Kth Largest, so the bounded heap flips to a max-heap (negated squared distance in Python): the root is the worst of the current best k, first out when a closer point arrives.

The Approach

For each point push (−(x² + y²), x, y); past size k, pop — evicting the farthest candidate. Survivors are the k closest. Squared distance preserves order because squaring is monotone on non-negatives; saying that sentence preempts the 'where is sqrt?' question.

Quickselect on squared distance gives O(n) average for the one-shot case, and heapq.nsmallest(k, points, key=...) is the stdlib one-liner — same trade conversation as Kth Largest, transplanted.

Python Solution

import heapq


def k_closest(points: list[list[int]], k: int) -> list[list[int]]:
    """K points nearest the origin via a bounded max-heap."""
    heap: list[tuple[int, int, int]] = []
    for x, y in points:
        heapq.heappush(heap, (-(x * x + y * y), x, y))
        if len(heap) > k:
            heapq.heappop(heap)        # drop the farthest candidate
    return [[x, y] for _, x, y in heap]

Complexity

• Time: O(n log k) — one bounded heap operation per point
• Space: O(k) — the candidate heap

Interview Tips and Follow-Ups

• Ties at the k-th distance: any valid subset passes — but confirm, because 'deterministic output' changes the tie handling.
• Negating only the distance (not the coordinates) in the tuple is the detail that keeps comparisons correct — tuple ordering compares fields in sequence.
• The geometric framing recurs at Meta as 'k nearest drivers/restaurants' — same code, plus a conversation about geo-indexing when n is planetary.

More Top-K and Heaps problems — and the other eight patterns — live in the Technical Interview category. Drill the pattern, not the problem: that is the entire thesis of this series.