K Closest Points to Origin

Why it belongs on this sheet

This is a high-yield heap problem and a very natural way to practice top-k thinking.

Pattern

Max-heap of size k

Approach

Keep only the best k points in a max-heap by distance. If a new point is closer than the farthest currently kept point, replace it.

Java solution

import java.util.PriorityQueue;

class Solution {
  public int[][] kClosest(int[][] points, int k) {
    PriorityQueue<int[]> heap = new PriorityQueue<>(
      (a, b) -> Integer.compare(distance(b), distance(a))
    );

    for (int[] point : points) {
      heap.offer(point);
      if (heap.size() > k) {
        heap.poll();
      }
    }

    int[][] answer = new int[k][2];
    for (int i = 0; i < k; i++) {
      answer[i] = heap.poll();
    }

    return answer;
  }

  private int distance(int[] point) {
    return point[0] * point[0] + point[1] * point[1];
  }
}

Complexity

• Time: O(n log k)
• Space: O(k)

Interview note

Squared distance is enough. Taking square roots only adds noise.