Find the k points nearest the origin. Keep a max-heap of k seats and bump whoever sits farthest, so the k closest are left at the end.
▼The problem
LeetCode 973 (Medium). Given an array of points [x, y] on the plane and an integer k, return the k points closest to the origin (0, 0), in any order (the answer is unique up to order).
Examples (LeetCode's): [[1,3],[-2,2]], k = 1 → [[-2,2]] (scores 10 and 8); [[3,3],[5,-1],[-2,4]], k = 2 → [[3,3],[-2,4]] (scores 18, 26, 20).
The solution
import heapq
def kClosest(points, k):
heap = [] # scores negated
for x, y in points:
heapq.heappush(heap, (-(x*x + y*y), x, y))
if len(heap) > k:
heapq.heappop(heap) # bump the farthest
return [[x, y] for _, x, y in heap]Transcript
K Closest Points to Origin. Given points on a grid and a number k, return the k points closest to the origin, zero zero.
Compare x squared plus y squared; the square root never changes the order, so skip it. One three scores ten, and minus two two scores eight, so for k equals one, it's minus two two. Three three, five minus one and minus two four score eighteen, twenty-six and twenty. For k equals two, keep three three and minus two four.
The simple way scores every point, sorts all n, and takes the first k. That's n log n time, sorting far more than you need.
Better: a campfire with a bench of just k seats. Each newcomer sits down, and if more than k are seated, whoever sits farthest from the fire is bumped. The bench is a max-heap, so the farthest is always on top. Whoever is left is the k closest.
Python's heap keeps the smallest on top, so push each point with its score negated. When the heap holds more than k, pop. Return what remains.
Six points, k equals three. Twenty, ten and thirty-four fill the bench. Five arrives: thirty-four is bumped. Seventeen arrives: twenty is bumped. Twenty-five arrives, but it's the farthest, so it's bumped at once. The bench keeps five, ten and seventeen.
Each push and pop costs log k, so it's n log k time and k space. Quickselect averages linear time, but the heap is simpler and works on a stream.
Square, don't root. Keep k seats, bump the farthest. That's K Closest Points to Origin.