堆(优先队列)的实现原理
#数据结构#堆 2023-10-12

堆(优先队列)是数据结构与算法中非常经典的结构，被广泛应用到计算机科学当中. 例如在堆排和优先级调度中堆便是核心的数据结构. 本文介绍堆的基本实现原理，力求清晰易懂.

1. 基础概念

In computer science, a heap is a specialized tree-based data structure that satisfies the heap property: In a max heap, for any given node C, if P is a parent node of C, then the key (the value) of P is greater than or equal to the key of C. In a min heap, the key of P is less than or equal to the key of C.

堆(优先队列)的这种结构是为了在堆进行操作时能够尽可能降低算法的复杂度，从使用角度讲，堆(优先队列)支持初始化、队头弹出、插入等操作。并保证从队头取出的元素具有最高的优先级。例如以元素大小作为优先级则可构造大(小)根堆，保证每次从队头取出的都是最大(小)的元素。

堆的实现一般是由数组模拟完全二叉树的数据结构实现。用数组存储的完全二叉树中，下标为 i 的节点的孩子节点(如果存在)下标分别为 2i + 1 和 2i + 2. 同理，节点 i 的父节点(如果存在)为 (i - 1) / 2。

定义优先队列 Heap 的 API，其中 cmp 用以比较两个元素的优先级：

1
2
3
4
5
6
7
8
9
10
11
12
13
class Heap<T> {
  q: T[];
  cmp: (a: T, b: T) => boolean;

  constructor(data: T[], cmp?: (a: T, b: T) => boolean) {}

  get size() {
    return this.q.length
  }

  pop(): T | undefined {}
  push(x: T) {}
}

2. 核心算法

2.1 节点下沉 (sink)

如果有一个二叉树，它的左子树和右子树都已经是大根堆，为了使该二叉树构成堆，需要把根节点“下沉 (sink)“ 到合适的位置：

  • 将该二叉树根节点的值与左右子树的根节点做对比，
  • 如果它的值大于左右子树根节点则此二叉树也构成一个大根堆. 算法结束.
  • 否则将它与比自身大的那个节点的位置进行交换
  • 此时继续考虑交换后节点所在的子树，递归执行上述操作

通过以上算法，初始的根节点被“下沉”到合适的位置，最终形成一个大根堆。

下沉算法的递归实现如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
sink(i: number) {
  let max = 2 * i + 1

  if (max >= this.size) return

  if (max + 1 < this.size && this.cmp(this.q[max + 1], this.q[max])) {
    ++max
  }

  if (this.cmp(this.q[i], this.q[max])) return;

  [this.q[i], this.q[max]] = [this.q[max], this.q[i]]
  this.sink(max)
}

实际上，下沉算法往往使用迭代实现：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
sink(i: number) {
  while (2 * i + 1 < this.size) {
    let max = 2 * i + 1

    if (max + 1 < this.size && this.cmp(this.q[max + 1], this.q[max])) {
      ++max
    }

    if (this.cmp(this.q[i], this.q[max])) break;

    [this.q[i], this.q[max]] = [this.q[max], this.q[i]]

    i = max
  }
}

2.2 节点上浮 (swim)

对于一个堆，如果在它的末尾插入一个元素，为了保证堆的性质需要把该节点“上浮 (swim)”到合适的位置：

  • 将该元素与它的父节点进行比较，如果满足堆的定义，则此时二叉树构成堆，算法结束
  • 否则将他与父元素交换
  • 递归执行上述操作

当算法结束时，新插入的节点上浮 (swim) 到合适的位置使堆的性质得以保证.

上浮算法的递归实现如下：

1
2
3
4
5
6
7
8
9
10
swim(i: number) {
  const p = Math.floor((i - 1) / 2)

  if (p < 0) return

  if (this.cmp(this.q[p], this.q[i])) return;
  [this.q[p], this.q[i]] = [this.q[i], this.q[p]]

  this.swim(p)
}

实际上，常用迭代实现上浮算法：

1
2
3
4
5
6
7
8
9
swim(i: number) {
  while (Math.floor((i - 1) / 2) >= 0) {
    const p = Math.floor((i - 1) / 2)
    if (this.cmp(this.q[p], this.q[i])) break;

    [this.q[p], this.q[i]] = [this.q[i], this.q[p]]
    i = p
  }
}

3. API 实现

有了上浮和下沉算法，堆的 API 实现起来就非常简单直观了.

3.1 插入元素

往一个堆里插入元素时，只需要把该元素放到堆的末尾，然后把该元素上浮到合适的位置即可：

1
2
3
4
push(x: T) {
  this.q.push(x)
  this.swim(this.size - 1)
}

3.2 弹出元素

由于堆的性质，队首的元素一定是优先级最高的元素. 因此直接把该元素弹出即可. 为了保证剩下数据仍然符合堆的性质，需要把末尾的元素放到队首，然后把它下沉到合适的位置即可.

1
2
3
4
5
6
7
8
pop(): T | undefined {
  const res = this.q[0]
  this.q[0] = this.q[this.size - 1]
  this.q.pop()
  this.sink(0)

  return res
}

3.3 初始化建堆

建堆是指对于一个空堆，给定一组元素把这些元素组织成一个堆. 一个简单的思路是一次把这些元素 push 到堆中. 但实际上往往使用一种更高效的方法, 这种方法在原有的数组上进行调整，最后使之成为堆.

对于给定的数组直接把它看成一个完全二叉树，显然这棵树的所有叶子节点都是一个堆. 那么对于这些叶子节点的父节点满足下沉算法的初始条件，则对这些父节点执行下沉算法后，新的堆被建立. 递归的考虑以这些堆的父节点为根节点的二叉树，对他们执行下沉算法，逐级完成堆的构建，直至抵达根节点.

简而言之，从数组中的最后一个非叶子节点倒序遍历该数组并执行下沉算法，即完成堆的建立:

1
2
3
4
5
init() {
  for (let i = Math.floor((this.size - 2) / 2); i >= 0; i--) {
    this.sink(i)
  }
}

3.4 完整实现

堆的实现主要依赖两个核心算法: 下沉和上浮，具体 API 的实现都是基于这两个算法的简单使用.

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
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
export default class Heap<T> {
  q: T[];
  cmp: (a: T, b: T) => boolean;

  constructor(data: T[], cmp?: (a: T, b: T) => boolean) {
    this.q = data
    this.cmp = cmp || ((a: T, b: T) => (a as number) - (b as number) > 0);
    this.init()
  }

  get size(): number {
    return this.q.length
  }

  init() {
    for (let i = Math.floor((this.size - 2) / 2); i >= 0; i--) {
      this.sink(i)
    }
  }

  sink(i: number) {
    while (2 * i + 1 < this.size) {
      let max = 2 * i + 1

      if (max + 1 < this.size && this.cmp(this.q[max + 1], this.q[max])) {
        ++max
      }

      if (this.cmp(this.q[i], this.q[max])) break;

      [this.q[i], this.q[max]] = [this.q[max], this.q[i]]

      i = max
    }
  }

  pop(): T | undefined {
    const res = this.q[0]
    this.q[0] = this.q[this.size - 1]
    this.q.pop()
    this.sink(0)

    return res
  }

  push(x: T) {
    this.q.push(x)
    this.swim(this.size - 1)
  }

  swim(i: number) {
    while (Math.floor((i - 1) / 2) >= 0) {
      const p = Math.floor((i - 1) / 2)
      if (this.cmp(this.q[p], this.q[i])) break;

      [this.q[p], this.q[i]] = [this.q[i], this.q[p]]
      i = p
    }
  }
}