K Closest Points to Origin

Signal

"Top-K by some derived distance/score, order among the K doesn't matter" is a heap-selection problem, not a full sort — you only need the k smallest distances, not every point ranked.

Approach

Compute squared Euclidean distance for each point (no need for sqrt — it's monotonic, so comparisons agree either way) and keep the k smallest by that key. A max-heap capped at size k works for a true streaming version; here, with all points available at once, heapq.nsmallest is the direct, correct tool for exactly this shape of query.

Skeleton

key = lambda p: p[0]**2 + p[1]**2
return nsmallest(k, points, key=key)

Solution

import heapq

def k_closest(points: list[list[int]], k: int) -> list[list[int]]:
    return heapq.nsmallest(k, points, key=lambda p: p[0] ** 2 + p[1] ** 2)

Complexity

O(n log k) time — nsmallest maintains an internal heap of size k while scanning all n points once. O(k) space for the heap (ignoring the O(n) input).

Pitfalls

• Computing sqrt for the distance is wasted work — squared distance preserves the same ordering and avoids floating point entirely.
• Sorting all n points by distance and slicing the first k works but is O(n log n); nsmallest's heap-based approach is O(n log k), which matters when k is much smaller than n.
• If an interviewer turns this into a true stream (points arrive one at a time, k fixed), switch to a max-heap of size k, popping the current largest whenever a closer point arrives — the same pattern as kth-largest-in-a-stream.