K Closest Points to Origin

Difficulty: Medium

You're given a list of points on a 2-D plane, each written as [x, y], and a number k. Find the k points that are closest to the origin (0, 0), using ordinary straight-line (Euclidean) distance.

Return those k points in any order - as long as the set of points is correct, the order doesn't matter.

Examples

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

Distance of (1,3) from the origin is sqrt(1+9) = sqrt(10) ~ 3.16. Distance of (-2,2) is sqrt(4+4) = sqrt(8) ~ 2.83, which is smaller, so it's the closer point.

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

Squared distances are 18, 26, and 20. The two smallest belong to (3,3) and (-2,4).

Constraints

  • 1 <= k <= points.length <= 10^4

  • -10^4 <= x[i], y[i] <= 10^4

Approach

Since you only need the k closest points, not a full ranking of every point, sorting everything is more work than necessary. Comparing squared distances (skipping the square root, which doesn't change the ordering) already saves some work either way.

The efficient approach keeps a max-heap of size k, ordered by distance. For each point, add it to the heap; once the heap grows past size k, remove the point with the largest distance - it's now guaranteed not to be among the k closest. Whatever remains in the heap at the end is the answer.

Solutions

Brute Force - Sort All Points

Compute every point's squared distance from the origin, sort all points by that distance, and take the first k.

function kClosest(points, k) {
  const sorted = [...points].sort((a, b) => {
    const distA = a[0] * a[0] + a[1] * a[1];
    const distB = b[0] * b[0] + b[1] * b[1];
    return distA - distB;
  });
  return sorted.slice(0, k);
}

Time: O(n log n) · Space: O(n) for the sorted copy

Optimal - Max-Heap of Size k

Keep a max-heap (ordered by distance) that never grows past size k. Adding a new point that makes the heap too big evicts the currently-farthest point.

class MaxHeap {
  constructor(compare) {
    this.data = [];
    this.compare = compare; // compare(a, b) > 0 means a should end up above b
  }

  size() {
    return this.data.length;
  }

  push(val) {
    this.data.push(val);
    let i = this.data.length - 1;
    while (i > 0) {
      const parent = (i - 1) >> 1;
      if (this.compare(this.data[i], this.data[parent]) > 0) {
        [this.data[i], this.data[parent]] = [this.data[parent], this.data[i]];
        i = parent;
      } else break;
    }
  }

  pop() {
    const top = this.data[0];
    const last = this.data.pop();
    if (this.data.length > 0) {
      this.data[0] = last;
      let i = 0;
      const n = this.data.length;
      while (true) {
        let largest = i;
        const left = 2 * i + 1;
        const right = 2 * i + 2;
        if (left < n && this.compare(this.data[left], this.data[largest]) > 0) largest = left;
        if (right < n && this.compare(this.data[right], this.data[largest]) > 0) largest = right;
        if (largest === i) break;
        [this.data[i], this.data[largest]] = [this.data[largest], this.data[i]];
        i = largest;
      }
    }
    return top;
  }
}

function kClosest(points, k) {
  const heap = new MaxHeap((a, b) => a.dist - b.dist);

  for (const point of points) {
    const dist = point[0] * point[0] + point[1] * point[1];
    heap.push({ point, dist });
    if (heap.size() > k) heap.pop();
  }

  const result = [];
  while (heap.size() > 0) result.push(heap.pop().point);
  return result;
}

Time: O(n log k) - each of the n points does an O(log k) heap operation · Space: O(k) - the heap never holds more than k points