Drone Nearest Drop-Offs

MediumheapAlso known as K Closest Points to Origin

A delivery drone lives on a rooftop pad. The town is mapped on a grid in metres with the pad at (0, 0), so a drop-off at (30, -40) is 30 metres east and 40 metres south. The drone flies in straight lines.

On a windy afternoon the operator only allows short hops, so the drone will make just k deliveries today: the k drop-offs closest to the pad, as the crow flies. The rest wait for tomorrow.

Your task

Write nearestDropOffs(points, k). points is a list of [x, y] drop-off positions. Return the k points with the smallest straight-line distance to (0, 0), in any order.

The input is chosen so the answer is unique: the k-th closest point is strictly closer than every point left out. The same position can appear twice (two parcels for one building); each counts as its own drop-off.

nearestDropOffs(points: int[][], k: int) → int[][] Any order of the returned list is accepted.

Examples

Example 1

Input

points = [[30, -40], [10, 20], [-5, 0], [60, 0]]
k = 2

Output

[[10, 20], [-5, 0]]

Why

Distances: (30, -40) is 50 m, (10, 20) is about 22.4 m, (-5, 0) is 5 m, (60, 0) is 60 m. The two closest are (-5, 0) and (10, 20).

Example 2

Input

points = [[3, 3], [0, 5], [-1, -1]]
k = 2

Output

[[3, 3], [-1, -1]]

Why

(3, 3) is about 4.24 m away and (0, 5) is 5 m, so (3, 3) is closer even though its coordinates add up to 6, more than 5. Adding up |x| + |y| measures the wrong distance. The answer is (-1, -1) and (3, 3).

Constraints

1 ≤ k ≤ points.length ≤ 10,000
-10,000 ≤ x, y ≤ 10,000
The k chosen points are strictly closer than all other points.

Start with brute force

Compute every distance, sort all n points by it, and take the first k. This is correct and takes O(n log n).

That is perfectly usable for 10,000 points. The heap version is O(n log k), which matters when k is small and the list is huge, or when points stream in and you cannot hold them all. The more common bug is not speed but measuring distance as |x| + |y|, or as x + y, which goes wrong for negative coordinates.

Hints

0 of 4 shown

Hint 1
To compare the distances of (3, 4) and (0, 6), do you actually need the square root? Which is bigger, 3² + 4² or 0² + 6²?
Hint 2
The square root keeps order, so compare x² + y² directly. That keeps everything in whole numbers. Now the task is: pick the k smallest of n numbers. Sorting all of them does more work than needed when k is small.
Hint 3
Walk through the points one by one and keep only the best k seen so far. When a new point arrives, which of your kept points is it competing against?
Hint 4
Keep a max-heap of size k, ordered by squared distance. Push each point; if the heap grows past k, pop the farthest. At the end the heap holds exactly the k nearest. In Python, push (-dist, x, y) onto heapq.

Walkthrough and solution

Unlocks when you pass every test. Stuck for a while? Looking is fine too, as long as you then close it and write the code yourself.

How to solve it

Two ideas. First, compare x² + y² instead of the true distance: the square root never changes which point is closer, and whole numbers avoid rounding. Second, keep only the best k so far in a max-heap, so the farthest of them is always on top, ready to be thrown out.

Trace the first example with k = 2. Squared distances: (30, -40) → 2500, (10, 20) → 500, (-5, 0) → 25, (60, 0) → 3600.

Push (30, -40). Heap: {2500}.
Push (10, 20). Heap: {2500, 500}. Size 2, fine.
Push (-5, 0). Size 3, so pop the farthest: 2500 goes. Heap: {500, 25}.
Push (60, 0). Size 3, pop the farthest: 3600 leaves straight away. Heap: {500, 25}.

The heap holds (10, 20) and (-5, 0). The result comes out in heap order, which is fine because any order is accepted.

Reference solution

function nearestDropOffs(points, k) {
  const d = ([x, y]) => x * x + y * y; // squared distance is enough
  const heap = []; // max-heap by distance: farthest kept point on top
  for (const p of points) {
    heap.push(p); // sift up
    for (let i = heap.length - 1; i > 0; ) {
      const par = (i - 1) >> 1;
      if (d(heap[par]) >= d(heap[i])) break;
      [heap[par], heap[i]] = [heap[i], heap[par]];
      i = par;
    }
    if (heap.length > k) { // drop the farthest: last to top, sift down
      heap[0] = heap.pop();
      for (let i = 0; ; ) {
        const l = 2 * i + 1, r = l + 1;
        let m = i;
        if (l < heap.length && d(heap[l]) > d(heap[m])) m = l;
        if (r < heap.length && d(heap[r]) > d(heap[m])) m = r;
        if (m === i) break;
        [heap[m], heap[i]] = [heap[i], heap[m]];
        i = m;
      }
    }
  }
  return heap;
}

import heapq
 
def nearest_drop_offs(points, k):
    heap = []  # entries (-squared distance, x, y): farthest kept point on top
    for x, y in points:
        heapq.heappush(heap, (-(x * x + y * y), x, y))
        if len(heap) > k:
            heapq.heappop(heap)  # drop the farthest
    return [[x, y] for _, x, y in heap]

TimeO(n log k)Each of the n points costs one push and at most one pop on a heap of size k + 1.

SpaceO(k)The heap never holds more than k + 1 points.

Learn the idea: Heaps and Priority Queues Math for Algorithms

Ctrl+Enter runs

Run your code to see results here. Run checks the examples. Submit checks every test, including hidden edge cases.

Write code