K Closest Points to Origin

medium array heap

Problem

Given an array of points on the plane and an integer k, return the k points closest to the origin (0, 0) by Euclidean distance.

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

Output[[-2,2],[0,1]]

Keep a max-heap of size k keyed by squared distance; pop when bigger arrives.

points (e.g. 1,3; -2,2)

import heapq
def k_closest(points, k):
    heap = []
    for x, y in points:
        d = -(x * x + y * y)  # negate for max-heap
        if len(heap) < k:
            heapq.heappush(heap, (d, x, y))
        elif d > heap[0][0]:
            heapq.heapreplace(heap, (d, x, y))
    return [[x, y] for _, x, y in heap]

// Sketch with a sorted array stand-in for a max-heap of size k.
function kClosest(points, k) {
  const heap = [];  // max-heap by sq distance
  for (const [x, y] of points) {
    const d = x * x + y * y;
    if (heap.length < k) {
      heap.push([d, x, y]);
      heap.sort((a, b) => b[0] - a[0]);
    } else if (d < heap[0][0]) {
      heap[0] = [d, x, y];
      heap.sort((a, b) => b[0] - a[0]);
    }
  }
  return heap.map(([_, x, y]) => [x, y]);
}

class Solution {
    public int[][] kClosest(int[][] points, int k) {
        PriorityQueue<int[]> heap = new PriorityQueue<>((a, b) -> b[0] - a[0]);
        for (int[] p : points) {
            int d = p[0] * p[0] + p[1] * p[1];
            if (heap.size() < k) heap.offer(new int[]{ d, p[0], p[1] });
            else if (d < heap.peek()[0]) { heap.poll(); heap.offer(new int[]{ d, p[0], p[1] }); }
        }
        int[][] out = new int[k][2];
        int i = 0;
        for (int[] e : heap) out[i++] = new int[]{ e[1], e[2] };
        return out;
    }
}

vector<vector<int>> kClosest(vector<vector<int>>& points, int k) {
    priority_queue<tuple<int, int, int>> heap;  // default = max-heap
    for (auto& p : points) {
        int d = p[0] * p[0] + p[1] * p[1];
        if ((int)heap.size() < k) heap.emplace(d, p[0], p[1]);
        else if (d < get<0>(heap.top())) { heap.pop(); heap.emplace(d, p[0], p[1]); }
    }
    vector<vector<int>> out;
    while (!heap.empty()) { auto [d, x, y] = heap.top(); heap.pop(); out.push_back({ x, y }); }
    return out;
}

Explanation

We want the k points nearest the origin. A neat way is to keep only the k best points seen so far in a max-heap keyed by distance, so the heap's top is always the worst of our current best.

To compare distances we use squared distance x*x + y*y instead of the real Euclidean distance. The square root does not change the ordering, so skipping it saves work and avoids floating point.

For each point: if the heap has fewer than k points, just add it. Otherwise, compare it with the heap's top (the farthest kept point). If the new point is closer, we evict the top and insert the new one; if not, we ignore it.

Example: points = [[1,3],[-2,2],[5,8],[0,1]], k = 2. Distances² are 10, 8, 89, 1. The heap fills with 10 and 8; then 89 is worse than the top 10, so skip; then 1 is closer than 10, so 10 is evicted. The heap ends with points of distance 8 and 1, i.e. [-2,2] and [0,1].

Because the heap never grows past size k, each insert costs only log k, making this faster than sorting all points when k is small.

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