Heaps and Priority Queues
Top-k is the most common heap question, and the counter-intuitive part is that you use a min-heap to find the largest items.
Keep a min-heap of size k. Its root is the smallest of your current best k. When a new number arrives, compare it to that root - if it is bigger, the root is no longer in the top k, so evict it.
K largest
javascript
// From the Heaps lesson, repeated so this example runs on its own.
class MinHeap {
items = []
get size() { return this.items.length }
push(v) {
this.items.push(v)
let i = this.items.length - 1
while (i > 0) {
const p = (i - 1) >> 1
if (this.items[p] <= this.items[i]) break
;[this.items[p], this.items[i]] = [this.items[i], this.items[p]]
i = p
}
}
pop() {
if (!this.items.length) return undefined
const top = this.items[0]
const last = this.items.pop()
if (this.items.length) {
this.items[0] = last
let i = 0
while (true) {
const l = 2 * i + 1, r = 2 * i + 2
let s = i
if (l < this.items.length && this.items[l] < this.items[s]) s = l
if (r < this.items.length && this.items[r] < this.items[s]) s = r
if (s === i) break
;[this.items[i], this.items[s]] = [this.items[s], this.items[i]]
i = s
}
}
return top
}
}
function kLargest(nums, k) {
const heap = new MinHeap()
for (const n of nums) {
heap.push(n)
// Never let the heap exceed k, so the root is always
// the weakest member of the current top k.
if (heap.size > k) heap.pop()
}
const out = []
while (heap.size) out.push(heap.pop())
return out.reverse()
}
console.log(kLargest([3, 1, 5, 12, 2, 11], 3)) // [12, 11, 5]O(n log k) time and O(k) space. Sorting is O(n log n) time and O(n) space. When k is small and n is huge - the usual case - the heap wins decisively.
Count, then heap
javascript
// From the Heaps lesson, repeated so this example runs on its own.
class PriorityQueue {
items = []
constructor(compare = (a, b) => a - b) { this.compare = compare }
get size() { return this.items.length }
push(v) {
this.items.push(v)
let i = this.items.length - 1
while (i > 0) {
const p = (i - 1) >> 1
if (this.compare(this.items[p], this.items[i]) <= 0) break
;[this.items[p], this.items[i]] = [this.items[i], this.items[p]]
i = p
}
}
pop() {
if (!this.items.length) return undefined
const top = this.items[0]
const last = this.items.pop()
if (this.items.length) {
this.items[0] = last
let i = 0
while (true) {
const l = 2 * i + 1, r = 2 * i + 2
let s = i
if (l < this.items.length && this.compare(this.items[l], this.items[s]) < 0) s = l
if (r < this.items.length && this.compare(this.items[r], this.items[s]) < 0) s = r
if (s === i) break
;[this.items[i], this.items[s]] = [this.items[s], this.items[i]]
i = s
}
}
return top
}
}
function topKFrequent(items, k) {
const counts = new Map()
for (const item of items) counts.set(item, (counts.get(item) || 0) + 1)
const heap = new PriorityQueue((a, b) => a.count - b.count)
for (const [value, count] of counts) {
heap.push({ value, count })
if (heap.size > k) heap.pop()
}
const out = []
while (heap.size) out.push(heap.pop().value)
return out.reverse()
}
console.log(topKFrequent(["a", "b", "a", "c", "a", "b"], 2)) // ["a", "b"]Closest by distance
javascript
function kClosest(points, k) {
// Compare squared distance — no need for the square root.
const heap = new PriorityQueue(
(a, b) => (b.x * b.x + b.y * b.y) - (a.x * a.x + a.y * a.y)
)
for (const point of points) {
heap.push(point)
if (heap.size > k) heap.pop()
}
const out = []
while (heap.size) out.push(heap.pop())
return out
}Skipping Math.sqrt is worth doing. It does not change the ordering and it removes a call from the hot path.
[...nums].sort((a,b)=>b-a).slice(0,k) is one line and obviously correct.The heap wins when n is large, k is small, or the data arrives as a stream you cannot sort because you never hold all of it.