Marauder Bot

Problem

You are given a list of points on a 2D plane and an integer k. Return the k points closest to the origin (0, 0), measured by straight-line distance. Any order of the answer is acceptable.

Example. With points [[1, 3], [-2, 2]] and k = 1, the distances from the origin are roughly 3.16 and 2.83, so the closer point [-2, 2] is the answer.

Key idea

Sorting every point by distance and taking the first k works, but it pays for a full ordering of all n points when only the k smallest distances actually matter. That extra work can be avoided by keeping just the k best candidates seen so far instead of ranking everything.

Maintain a max-heap capped at size k, ordered by squared distance from the origin to avoid unnecessary square roots. Process each point: add it to the heap, and if the heap now exceeds k elements, remove the one with the largest distance, which sits at the top. After processing every point, the heap contains exactly the k closest, because any point farther than the current worst of the top k gets evicted as soon as a closer candidate displaces it.

An alternative for very large inputs is a quickselect-style partition on distance, which can average better than the heap's O(n log k) but has worse worst-case behavior and no streaming-friendly property.

Solution

// Minimal binary heap; the comparator (hasPriority) decides min-heap or max-heap order.
class Heap<T> {
  private readonly items: T[] = [];

  constructor(private readonly hasPriority: (a: T, b: T) => boolean) {}

  get size(): number {
    return this.items.length;
  }

  push(item: T): void {
    this.items.push(item);
    let i = this.items.length - 1;
    while (i > 0) {
      const parent = (i - 1) >> 1;
      if (this.hasPriority(this.items[parent], this.items[i])) {
        break;
      }
      [this.items[parent], this.items[i]] = [this.items[i], this.items[parent]];
      i = parent;
    }
  }

  pop(): T {
    const top = this.items[0];
    const last = this.items.pop()!;
    if (this.items.length > 0) {
      this.items[0] = last;
      this.bubbleDown(0);
    }
    return top;
  }

  private bubbleDown(start: number): void {
    let i = start;
    while (true) {
      const left = 2 * i + 1;
      const right = 2 * i + 2;
      let next = i;
      if (left < this.items.length && this.hasPriority(this.items[left], this.items[next])) {
        next = left;
      }
      if (right < this.items.length && this.hasPriority(this.items[right], this.items[next])) {
        next = right;
      }
      if (next === i) {
        break;
      }
      [this.items[i], this.items[next]] = [this.items[next], this.items[i]];
      i = next;
    }
  }
}

function squaredDistance(point: number[]): number {
  return point[0] * point[0] + point[1] * point[1]; // avoids the sqrt; only relative order matters
}

function kClosest(points: number[][], k: number): number[][] {
  const heap = new Heap<number[]>((a, b) => squaredDistance(a) > squaredDistance(b)); // max-heap by distance

  for (const point of points) {
    heap.push(point);
    if (heap.size > k) {
      heap.pop(); // evict the current farthest point once over k
    }
  }

  const result: number[][] = [];
  while (heap.size > 0) {
    result.push(heap.pop());
  }

  return result;
}

Complexity

Time: O(n log k). Each of the n points does an O(log k) heap operation.
Space: O(k). The heap holds at most k points; O(1) extra if the answer can overwrite the input.

Watch out for

Compare squared distances, not raw coordinates or distances needing a square root on every comparison; it is faster and avoids floating-point noise.
Use a max-heap so the worst of the current top k is the one evicted, not a min-heap, which would evict the best.

Pattern

This is another instance of the bounded top-k heap pattern: keep only the k best candidates by a comparable key and let the heap evict the worst as better ones arrive, the same technique used for streaming kth-largest problems.

K Closest Points to Origin