K Closest Points to Origin

The problem

Each entry of points is a point [x, y] in the plane. Return the k points that are closest to the origin, the point [0, 0], in any order. The distance of a point from the origin is the square root of x × x + y × y. If several points are equally far away, any of them may fill the last places. A point that is listed twice counts as two points.

What to write

function solve(points: number[][], k: number): number[][]

Input
points = [[1,3],[-2,2]]
k = 1

Output
[[-2,2]]

Input
points = [[3,3],[5,-1],[-2,4]]
k = 2

Output
[[3,3],[-2,4]]

Constraints

1 ≤ k ≤ points.length.

Hints

Each hint gives away a little more. Open one only when you need it.

1. Hint 1: Use the constraints

   Any order is accepted, and when distances tie, any of the tied points will do. So you never have to report a distance. You only have to decide, for two points, which one is closer.

2. Hint 2: Ask a smaller question

   Which is closer to [0, 0]: [1, 3] or [-2, 2]? Work it out on paper. At which moment did you know the answer: after taking the square roots, or already before?

3. Hint 3: Name the pattern

   You want the k best of n items, where best means the smallest x × x + y × y. That request signals the pattern called heap of the best k: keep only the k closest points seen so far, in a heap whose top is the farthest of them, the first one to throw out. Read the lesson on heap of the best k.

4. Hint 4: State the invariant

   After each point, the heap holds the k closest points seen so far, with the farthest of them on top. A new point pushes that one out exactly when the heap has grown past k and the top is now the farthest of k + 1.

5. Hint 5: Sketch the pseudocode

   For each point compute x × x + y × y. Push the point onto a max-heap ordered by that number. If the heap now holds more than k points, pop its top, the farthest. When every point has been read, return the points left in the heap.

Solution

From here on the page gives the method away. If you are still solving, go back to the hints.

Approach

Compare squared distances to avoid square roots. Keep a max-heap of the k closest points so far, and remove the farthest whenever it holds more than k.

The slow way

Find the closest point by checking every point, move it to the answer, and repeat until you have k points. It is correct, but each of the k searches walks all the points that are left, so the work is about k × n steps, far more than needed when k is large.

The insight

A larger number always has a larger square root, so x × x + y × y already tells you which of two points is closer: no square root is needed. And you never need all the points in order, only the k closest. A heap of size k keeps exactly those. It always knows the farthest of them, the one a closer point replaces.

Step by step

1. Measure each point by its squared distance from the origin. The squared distance, x × x + y × y, orders the points exactly as the real distance does. It skips the square root and stays a whole number when the coordinates are whole.
2. Keep a max-heap of the k closest points seen so far. A max-heap always has its largest entry on top, here the farthest of the points it holds. That is the one point of the k that a closer newcomer should replace.
3. Put the next point into the heap. Every point gets its chance to be among the k closest. A push takes about log k steps, because the heap never holds more than k + 1 points.
4. If the heap now holds more than k points, remove the farthest. With k + 1 points in the heap, the farthest of them cannot be among the k closest overall: k points are at least as close. Removing it loses nothing.
5. The k points left in the heap are the k closest. Every point was put in, and a point was removed only when k others were at least as close. Any order is accepted, so the heap's own order will do.

The same numbered steps are the comments in each solution.

function solve(points, k) {
  // 1. Measure each point by its squared distance from the origin
  const squaredDistance = ([x, y]) => x * x + y * y;
  // JavaScript has no built-in heap, so these two helpers keep an array of points
  // in max-heap order by hand: no point is farther than its parent, so the
  // farthest is at index 0. The parent of index i is at (i - 1) >>> 1, which
  // is (i - 1) / 2 rounded down, and its two children are at 2 * i + 1 and
  // 2 * i + 2.
  function heapPush(heap, point) {
    // The new point starts at the end and swaps with its parent while the
    // parent is closer.
    heap.push(point);
    let i = heap.length - 1;
    while (i) {
      const p = (i - 1) >>> 1;
      if (squaredDistance(heap[p]) >= squaredDistance(heap[i])) break;
      [heap[p], heap[i]] = [heap[i], heap[p]];
      i = p;
    }
  }
  function heapPop(heap) {
    // The farthest point leaves. The last point takes its place at index 0,
    // then swaps with its farther child while that child is farther.
    const last = heap.pop();
    if (!heap.length) return;
    heap[0] = last;
    let i = 0;
    while (true) {
      let child = 2 * i + 1;
      if (child >= heap.length) break;
      if (
        child + 1 < heap.length &&
        squaredDistance(heap[child + 1]) > squaredDistance(heap[child])
      )
        child++;
      if (squaredDistance(heap[i]) >= squaredDistance(heap[child])) break;
      [heap[i], heap[child]] = [heap[child], heap[i]];
      i = child;
    }
  }
  // 2. Keep a max-heap of the k closest points seen so far
  const heap = [];
  for (const point of points) {
    // 3. Put the next point into the heap
    heapPush(heap, point);
    // 4. If the heap now holds more than k points, remove the farthest
    if (heap.length > k) heapPop(heap);
  }
  // 5. The k points left in the heap are the k closest
  return heap;
}

function solve(points: number[][], k: number): number[][] {
  // 1. Measure each point by its squared distance from the origin
  const squaredDistance = ([x, y]: number[]) => x * x + y * y;
  // TypeScript has no built-in heap, so these two helpers keep an array of points
  // in max-heap order by hand: no point is farther than its parent, so the
  // farthest is at index 0. The parent of index i is at (i - 1) >>> 1, which
  // is (i - 1) / 2 rounded down, and its two children are at 2 * i + 1 and
  // 2 * i + 2.
  function heapPush(heap: number[][], point: number[]): void {
    // The new point starts at the end and swaps with its parent while the
    // parent is closer.
    heap.push(point);
    let i = heap.length - 1;
    while (i) {
      const p = (i - 1) >>> 1;
      if (squaredDistance(heap[p]) >= squaredDistance(heap[i])) break;
      [heap[p], heap[i]] = [heap[i], heap[p]];
      i = p;
    }
  }
  function heapPop(heap: number[][]): void {
    // The farthest point leaves. The last point takes its place at index 0,
    // then swaps with its farther child while that child is farther.
    const last = heap.pop()!;
    if (!heap.length) return;
    heap[0] = last;
    let i = 0;
    while (true) {
      let child = 2 * i + 1;
      if (child >= heap.length) break;
      if (
        child + 1 < heap.length &&
        squaredDistance(heap[child + 1]) > squaredDistance(heap[child])
      )
        child++;
      if (squaredDistance(heap[i]) >= squaredDistance(heap[child])) break;
      [heap[i], heap[child]] = [heap[child], heap[i]];
      i = child;
    }
  }
  // 2. Keep a max-heap of the k closest points seen so far
  const heap: number[][] = [];
  for (const point of points) {
    // 3. Put the next point into the heap
    heapPush(heap, point);
    // 4. If the heap now holds more than k points, remove the farthest
    if (heap.length > k) heapPop(heap);
  }
  // 5. The k points left in the heap are the k closest
  return heap;
}

import heapq


def solve(points, k):
    # 1. Measure each point by its squared distance from the origin
    def squared_distance(point):
        x, y = point
        return x * x + y * y

    # 2. Keep a max-heap of the k closest points seen so far
    # heapq is a min-heap: heap[0] is its smallest entry. Each point goes in with
    # its distance negated, so the farthest point is the one at heap[0].
    heap = []
    for point in points:
        # 3. Put the next point into the heap
        heapq.heappush(heap, (-squared_distance(point), point))
        # 4. If the heap now holds more than k points, remove the farthest
        if len(heap) > k:
            heapq.heappop(heap)
    # 5. The k points left in the heap are the k closest
    return [point for _, point in heap]

Example 2, traced

Input
points = [[3,3],[5,-1],[-2,4]]
k = 2

Output
[[3,3],[-2,4]]

Why it works

After each point, the heap holds the k closest points seen so far. A point is removed only when the heap holds k + 1 points and it is the farthest of them, so k points at least as close remain. Once every point has been read, the k points left are the k closest of all.

Time and space

Time O(n log k), extra space O(k). Each of the n points is pushed once, and at most one point is popped after it. A push or a pop on a heap of at most k + 1 points takes about log k steps: O(n log k) time. The heap never holds more: O(k) space.

n and m are input sizes, h is a tree's height, V and E are a graph's vertices and edges. The bounds describe this reference solution.

Watch out

• Do not rank by |x| + |y|, the coordinates added without their signs. In example 2 that gives 6 for all three points, and [5, -1], the farthest point, stays in the answer.
• Use a max-heap, with the farthest point on top. A min-heap removes the closest point each time: example 1 returns [1, 3] instead of [-2, 2].
• Remove a point only when the heap holds more than k. With k or more, the heap never keeps k points: example 2 returns one point instead of two.

Pattern

Heap of the best k. Builds on Top K Frequent Elements.

Keep going