Problem overview
The method kClosest(points, k) receives an array of points, where points[i] = [x, y], and a number k. It returns the k points closest to the origin (0, 0), as an array of [x, y] pairs.
The result must be sorted: the closest point first, then points at equal distance by x, then by y. The answer is guaranteed to be unique, so every point left out is strictly farther from the origin than every point in the result.
For points = [[1, 3], [-2, 2]] and k = 1 the answer is [[-2, 2]]: the squared distances are 10 and 8. For points = [[3, 3], [5, -1], [-2, 4]] and k = 2 the squared distances are 18, 26 and 20, so the answer is [[3, 3], [-2, 4]]. For points = [[0, 2], [2, 0], [1, 1], [0, -2], [3, 3]] and k = 4 the answer is [[1, 1], [0, -2], [0, 2], [2, 0]]: three points are exactly 2 away, so they are ordered by x and then by y, and [3, 3] is left out.
The statement asks for O(n log k) time and O(k) extra space, where n is the number of points. Sorting all the points would take O(n log n).
Approach: keep only the k best in a max-heap
Two observations make this simple.
No square roots. The distance is the square root of x * x + y * y, and the square root grows with its argument. A point is closer than another exactly when its squared distance is smaller. The squared distance is a whole number, so comparisons are exact and cheap.
Only the k best matter. Instead of sorting everything, keep a collection of the k closest points seen so far, and make it easy to find the farthest of them. A max-heap does that: its top is always the biggest entry, and inserting or removing the top costs O(log k). For each new point, insert it. If the heap now holds k + 1 points, the farthest of them cannot be among the k closest, so remove the top. After the last point, the heap holds the answer.
A heap is only partly ordered, so removing its entries gives them from farthest to closest, and with ties in no particular order. The final step sorts the k entries by distance, then x, then y, which costs O(k log k).
The heap comes from a library that the lesson provides in its starter code. Its heaps, algorithms.DataStructures.Heap.MaxHeap and MinHeap, take an optional compare function like the one sort takes, and have insert(value), extract(), isEmpty() and n, the number of values in the heap. JavaScript has no heap of its own.
Step by step
- Create the heap. It is a
MaxHeapwith the compare function(a, b) => a[0] - b[0]. Each entry will be[dist, x, y], and the function orders entries by their first field, the distance. - Visit every point with a
for...ofloop. The patternlet [x, y]takes each point apart into its two numbers. - Compute the squared distance:
x * x + y * y. - Insert
[dist, x, y]into the heap. The entry keeps the point next to the number the heap orders by. - Trim the heap. If
heap.n > k, callheap.extract(), which removes the top: the farthest of thek + 1points. After each point, the heap holds thekclosest seen so far. - Prepare an empty array
closestto take the survivors out of the heap. - Loop while
!heap.isEmpty(). Each turn removes one entry, so there arekturns. - Push
heap.extract()ontoclosest.extract()removes the top entry and returns it. - Sort
closestwith a compare function that subtracts the distances,a[0] - b[0], and when that is 0 moves on with||to the x values (index 1) and then the y values (index 2). - Return the points without their distances.
closest.maptakes each entry apart as[dist, x, y]and returns[x, y].
The solution in JavaScript
var algorithms = require('algorithms')
class Solution {
kClosest(points, k) {
let heap = new algorithms.DataStructures.Heap.MaxHeap((a, b) => a[0] - b[0]);
for (let [x, y] of points) {
let dist = x * x + y * y;
heap.insert([dist, x, y]);
if (heap.n > k) {
heap.extract();
}
}
let closest = [];
while (!heap.isEmpty()) {
closest.push(heap.extract());
}
closest.sort((a, b) => a[0] - b[0] || a[1] - b[1] || a[2] - b[2]);
return closest.map(([dist, x, y]) => [x, y]);
}
}
The JavaScript details worth a closer look:
- Destructuring.
let [x, y] of pointsunpacks each two-element array into two variables. The same pattern appears in the parameter list of themapcallback,([dist, x, y]) => [x, y]. The namedistis unused there, which is fine: it only marks the first slot. - Chaining comparisons with
||. A compare function must return a number, and0means "equal". Since0is falsy,a[0] - b[0] || a[1] - b[1]falls through to the second difference only when the first is 0. The same trick adds the third level. - Sorting needs a compare function. Without one,
sort()compares elements as strings:[10, 9, 1].sort()gives[1, 10, 9], while[10, 9, 1].sort((a, b) => a - b)gives[1, 9, 10]. - The subtraction is safe. Coordinates are at most 10^4 in size, so a squared distance is at most 2 · 10^8, and JavaScript numbers hold such values exactly. Subtracting two of them can't misbehave.
- The compare function is the heap's ordering. A
MaxHeapkeeps the entry that the compare function ranks highest on top, so it is the farthest point that gets extracted. AMinHeapwith the same function would drop the closest point instead, which is the opposite of what is needed.
Tracing points = [[3, 3], [5, -1], [-2, 4]], k = 2
| point | entry inserted | heap size | extracted |
|---|---|---|---|
| [3, 3] | [18, 3, 3] | 1 | nothing |
| [5, -1] | [26, 5, -1] | 2 | nothing |
| [-2, 4] | [20, -2, 4] | 3 | [26, 5, -1] |
The heap now holds the entries for [3, 3] and [-2, 4]. The sort puts distance 18 before distance 20, and the final map returns [[3, 3], [-2, 4]].
Complexity
Time: O(n log k). Each of the n points costs one insertion and at most one extraction, and the heap never holds more than k + 1 entries, so each operation costs O(log k). Emptying the heap and sorting the survivors adds O(k log k), which is no more than O(n log k) because k is at most n.
Space: O(k). The heap holds at most k + 1 entries, and closest holds k.
Edge cases and common mistakes
- One point.
[[1, 1]]withk = 1returns[[1, 1]]. kequal to the number of points.[[3, 3], [5, -1], [-2, 4]]withk = 3returns[[3, 3], [-2, 4], [5, -1]]: nothing is dropped, and the sort orders them.- Repeated points.
[[1, 1], [1, 1], [2, 2]]withk = 2returns[[1, 1], [1, 1]]. The same point can appear more than once, and each copy is a separate entry. - Ties on distance. Several points can share a distance, as the third example shows. The heap doesn't order them, so the final sort must break ties by
xand theny. - Using the wrong kind of heap. A min-heap would keep the
kfarthest points. The heap that gets trimmed from the top must be a max-heap. - Forgetting to trim. Without the
heap.n > kcheck the heap keeps every point, and the algorithm turns into a slower full sort. - Taking a square root. It adds nothing and gives non-integer values for no gain.
- Returning the entries as they are. The entries are
[dist, x, y], and the answer wants[x, y].
Testing
points = [[1, 3], [-2, 2]],k = 1gives[[-2, 2]].points = [[3, 3], [5, -1], [-2, 4]],k = 2gives[[3, 3], [-2, 4]].points = [[0, 2], [2, 0], [1, 1], [0, -2], [3, 3]],k = 4gives[[1, 1], [0, -2], [0, 2], [2, 0]].points = [[1, 1]],k = 1gives[[1, 1]].points = [[1, 1], [1, 1], [2, 2]],k = 2gives[[1, 1], [1, 1]].
