Heaps and Priority Queues
Top K Problems
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.
Why a min-heap for the largest
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.
Top k most frequent
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"]K closest to a point
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
}
// A comparator-driven priority queue, as used above.
class PriorityQueue {
constructor(compare) { this.compare = compare; this.data = [] }
get size() { return this.data.length }
push(v) { this.data.push(v); this.data.sort(this.compare) }
pop() { return this.data.shift() }
peek() { return this.data[0] }
}
console.log(kClosest([{ x: 1, y: 1 }, { x: 5, y: 5 }, { x: 2, y: 2 }], 2))Skipping Math.sqrt is worth doing. It does not change the ordering and it removes a call from the hot path.
When sorting is better
- k is close to n - a heap of size n has no advantage over a sort.
- You need everything ordered, not just the top k.
- n is small. Under a few hundred items, sort and slice is clearer and just as fast.
- You need it once and clarity matters -
[...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.
Which heap to use is counter-intuitive
For the k largest elements you want a min-heap of size k, not a max-heap. The root is then the weakest element currently in your set, so each new candidate needs a single comparison against it: bigger means evict the root and insert, smaller means discard immediately.
K largest in O(n log k)
javascript
function topK(items, k) {
const heap = new MinHeap() // from the heaps lesson
for (const value of items) {
if (heap.size < k) {
heap.push(value)
} else if (value > heap.data[0]) {
heap.pop()
heap.push(value)
}
}
return heap.data
}
// MinHeap from the heaps lesson.
class MinHeap {
constructor() { this.data = [] }
get size() { return this.data.length }
push(v) {
this.data.push(v)
let i = this.data.length - 1
while (i > 0) {
const p = (i - 1) >> 1
if (this.data[p] <= this.data[i]) break
;[this.data[p], this.data[i]] = [this.data[i], this.data[p]]
i = p
}
}
pop() {
if (this.data.length <= 1) return this.data.pop()
const topValue = this.data[0]
this.data[0] = this.data.pop()
let i = 0
for (;;) {
const l = 2 * i + 1, r = l + 1
let small = i
if (l < this.data.length && this.data[l] < this.data[small]) small = l
if (r < this.data.length && this.data[r] < this.data[small]) small = r
if (small === i) break
;[this.data[small], this.data[i]] = [this.data[i], this.data[small]]
i = small
}
return topValue
}
}
console.log(topK([1, 9, 5, 3, 7], 3).sort((a, b) => a - b)) // [5, 7, 9]When sorting is the better answer
O(n log k) only beats O(n log n) when k is genuinely much smaller than n. For k close to n the heap does more bookkeeping than a sort, and for a one-off query on a modest array sort().slice(0, k) is simpler and fast enough.
The heap earns its place when the data arrives as a stream you cannot hold in memory, or when k is small and n is very large. Say which of those you are assuming.