~/DHRUVUpskilling
← board/DSA/Top K Elements/dsa-top-k-elements-04
Solved·10 Sept

K Closest Points to Origin

DifficultyMedium
PatternTop K Elements
TrackDSA
Snippet
tl;dr

You are given an array of points on a 2D plane. Each point has an (x, y) coordinate.

full write-up

K Closest Points to the Origin

Problem Statement

You are given an array of points on a 2D plane. Each point has an (x, y) coordinate.

Find the k points that are closest to the origin, (0, 0).

Note: The distance between two points on a plane is the Euclidean distance: √(x² + y²).

Constraints

  • 1 ≤ k ≤ points.length ≤ 10³
  • −10⁴ < x[i], y[i] < 10⁴

Solution

We use a max heap to keep track of the k closest points found so far.

Steps

  • We calculate the distance from the origin for the first k points, and push them into a max heap. The heap keeps the point with the largest distance at the top.
  • For every remaining point, we compare its distance to the distance of the point at the top of the heap (the current farthest point among our k closest):
    • If the new point is closer than the one at the top, it means it deserves to be in our top k. We remove the farthest point from the heap, and add this new, closer point instead.
    • If the new point is farther away, we simply ignore it, since it's not closer than what we already have.
  • After going through all the points, the heap contains exactly the k closest points. We pop them all out into our result array.

Code

function kClosest(points, k) {
    let maxHeap = new MaxHeap(),
        result = [];

    for (var i = 0; i < k; i++) {
        maxHeap.offer([points[i].distanceFromOrigin(), points[i]]);
    }

    for (var i = k; i < points.length; i++) {
        if (
            points[i].distanceFromOrigin() <
            maxHeap.peek()[1].distanceFromOrigin()
        ) {
            maxHeap.poll();
            maxHeap.offer([points[i].distanceFromOrigin(), points[i]]);
        }
    }

    for (var i = 0; i < k; i++) {
        let point = maxHeap.poll()[1];
        result.push(point);
    }
    return result;
}