K Closest Points to Origin

Problem Statement

Given an array of 2-D points points[i] = [x, y] and integer k, return the k points closest to the origin by Euclidean distance. Any order is acceptable.

Input

An array of integer [x, y] pairs and the count k.

points = [[1, 3], [-2, 2], [5, 8], [0, 1]]
k      = 2

Output

The k closest points to the origin.

[[0, 1], [-2, 2]]
 
// squared distances: 10, 8, 89, 1
// two smallest: 1 and 8

Constraints

• 1 ≤ k ≤ points.length ≤ 10⁴
• -10⁴ ≤ x, y ≤ 10⁴

Intuition

Identical structure to “K Closest Elements” — just 2-D Euclidean distance. Never compute the square root: comparisons are preserved by squared distance, and we avoid a floating-point op per element. Use a max-heap of size K keyed on x² + y²; evict the worst when we exceed K.

point      x²+y²   heap (max by dist)                 action
[1,3]       10    [(10, [1,3])]                       push
[-2,2]       8    [(10, [1,3]), (8, [-2,2])]          push (size == k)
[5,8]       89    push → size 3 → pop (89, [5,8])     evict (89 worst)
[0,1]        1    push → size 3 → pop (10, [1,3])     evict (10 worst)
 
heap now holds [(8, [-2,2]), (1, [0,1])] → answer

Solution

import heapq
 
def kClosest(points: list[list[int]], k: int) -> list[list[int]]:
    max_heap: list[tuple[int, list[int]]] = []     # (-distance², point)
    for x, y in points:
        dist_sq = x * x + y * y
        heapq.heappush(max_heap, (-dist_sq, [x, y]))
        if len(max_heap) > k:
            heapq.heappop(max_heap)                # evict the farthest
    return [p for _, p in max_heap]

function kClosest(points: number[][], k: number): number[][] {
  // max-heap: larger distance² → higher priority to be popped
  const h = new Heap<[number, number[]]>((a, b) => b[0] - a[0]);
  for (const [x, y] of points) {
    const d = x * x + y * y;
    h.push([d, [x, y]]);
    if (h.size > k) h.pop();
  }
  const out: number[][] = [];
  while (h.size > 0) out.push(h.pop()![1]);
  return out;
}

Complexity

Edge Case Thinking