K Closest Points to Origin

easy ~15 min

Write k_closest(points, k) that returns the k points nearest to (0, 0) by straight-line distance.

points is a list of [x, y] pairs with -10^4 <= x, y <= 10^4; 1 <= k <= len(points) <= 10^5.
Return them in any order.
If there's a tie at the cutoff (several points at the same distance as the k-th closest), any choice among the tied points is accepted.

Aim for O(n log n) or better.

k_closest([[3, 3], [-1, 2], [4, -5]], 2)   # [[-1, 2], [3, 3]]  in any order
k_closest([[0, 5], [5, 0]], 1)             # [[0, 5]] or [[5, 0]]

Show hint

comparing x*x + y*y gives the same order as the true distance, with no sqrt. Sorting by it works; a heap of size k or quickselect can do better than O(n log n).

Topic: Sorting, custom keys, coordinate compression. sorted(key=...), multi-key and stable sorts, cmp_to_key, compressing coordinates.

Read the visual guide