源码

来自半年后的说明：我把代码转放到 gists 里了，当时代码风格受 Python 影响较大，这里就不再更改了，毕竟大家本地都有自己的格式化工具。

前往Github获取本文源码。

介绍

通常情况下，堆指的是二叉堆，它是一颗完全二叉树。完全二叉树指的是要么是满二叉树（都填满了），要么最底层从左向右排列。这里给出一个例子：

一颗完全二叉树，同时也是二叉堆

二叉堆除了需要满足是一个完全二叉树之外，还必须满足下方的数据永远比上方的大（或小），也被称为堆序性质。因此，进行插入/删除操作的时候可能要破坏这个性质，所以就需要我们对这个堆进行重新调整。

由于堆序性质，我们可以很方便地在一个堆中求最小（或最大）值，所以它在需要动态插入数据并且求出最值的时候就显得非常有用了。

实现

由于堆是一个完全二叉树，那么我们就不用非得用链表的形式来实现了，用数组足以应对，我们还用上面的图为例：

用数组的形式表示堆

为了计算方便，我们这里不用0号元素，让数组下表从1开始。那么，对于任意节点array[i]，它的左节点就是array[i * 2]，右节点就是array[i * 2 + 1]，父节点就是array[Math.floor(i / 2)]。

整体结构

我们这里来实现一个最小堆，就是根节点是最小值的堆，定义如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
class Heap {
    vals = [ NaN ]

	get length() {
        return this.vals.length - 1
    }

    get top() {
        return this.vals[1] ?? NaN
    }

    pop() {}

    push(val) {}
}

其中，top这个getter使用了双问号运算符来判断this.vals[1]是否为null或者undefined，如果是则返回一个NaN。

插入

由于插入可能会破坏堆序性质，所以我们需要进行上滤（percolate up）操作，使得它能不断在一个堆中上升到合适的位置。用一个图来表示我们的意思：

上滤

有两个时机来停止这个操作：当元素已经到达顶端时，或者元素已经符合堆序性质时。代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
push(val) {
    this.vals.push(val)
    this._percolateUp()
}

_percolateUp() {
    let cur = this.vals.length - 1
    while (true) {
        const parent = Math.floor(cur / 2)
        // parent < 1 就是元素已经到达顶端
        // vals[cur] > vals[parent] 就满足了最小堆的要求：根节点最小。
        if (parent < 1 || this.vals[cur] > this.vals[parent]) {
            return
        }
        this._swap(parent, cur)
        cur = parent
    }
}

_swap(i, j) {
    const temp = this.vals[i]
    this.vals[i] = this.vals[j]
    this.vals[j] = temp
}

删除

删除指的是删除顶部的元素，在当前这个最小堆中，就是删除最小值。所以与插入相反，我们在删除时要进行下滤（percolate down）操作。用一个图表示我们的意思：

下滤

同样有两个时机停止这个操作：当元素已经到达底端（没有子节点）时，或元素满足堆序性质时。

不同于上滤，因为一个节点可能有多个子节点，当选择时我们就要尽可能挑最小的换上来，保证堆序性质。代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
pop() {
    // 先把队尾的元素换到第一个，然后把它弹出，再进行下滤操作
    this._swap(1, this.vals.length - 1)
    const val = this.vals.pop()
    this._percolateDown()
    
    // 最后返回被删除的最小值。
    return val
}

_percolateDown() {
    let cur = 1
    while (true) {
        const left = cur * 2
        const right = left + 1
        let child = left

        // 当前节点没有左节点了，那么它也不可能有右节点，也就是说已经到底了。
        if (left >= this.vals.length) {
            return
        }

        // 当 right >= vals.length 的时候，当前节点就没有右节点了。
        // 所以当右节点存在，并且左节点比右节点大的时候，我们把 child 重新赋值为 right，
        // 这样就能选出较小的子节点了。
        if (right < this.vals.length && this.vals[left] > this.vals[right]) {
            child = right
        }

        // 如果子节点比当前节点更大，满足了堆序性质，就停止操作。
        if (this.vals[child] > this.vals[cur]) {
            return
        }

        this._swap(child, cur)
        cur = child
    }
}

这样，我们的最小堆就基本实现完毕了。接下来进行一个实际应用，求最大的k个元素。

实际应用

对于求最大的k个元素，我们可以维护一个最小堆：如果堆中元素的数量还不到k个，那就直接把它加入堆中；否则，如果当前值比堆中的最小值大，那么就弹出堆的最小值，并且把当前值放入堆中。代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
function getTopK(nums, k) {
    const heap = new Heap()
    nums.forEach(n => {
        if (heap.length < k) {
            heap.push(n)
        } else if (n > heap.top) {
            heap.pop()
            heap.push(n)
        }
    })
    
    // 这一步只是把结果放到数组中返回
    const result = []
    while (heap.length > 0) {
        result.push(heap.pop())
    }
    return result
}

const top3 = getTopK([5, 4, 10, 1, 6, 9, 2, 8], 3)
console.log(top3)

我们来看一下整个复杂度：由于遍历了n个数字，此时的复杂度是O(n)；每一次遍历中我们都操作了堆，由于堆最多有k个元素，所以循环内部的复杂度是O(log k)。总体的复杂度是O(n*log k)。

我们把结果放到数组中时，由于有k个数，此时复杂度时O(k)；每一次循环都操作了堆，由于堆最多k个元素，所以循环内部的复杂度是O(log k)。总体的复杂度是O(k*log k)。

那我们知道，k <= n，所以加起来的复杂度还是O(n*log k)。

如果我们用排序的话，最少也要O(n*log n)，不如堆的操作快。