# 树和数组的千层套路

今天教你怎么在线性数组和二叉树间反复横跳。

二叉堆
树状数组
线段树

在最简单的一类数据结构中，有两种最为基础：线性数组和树，它们在实际的应用中十分常用，比如线性数组用来存储结构化数据，树用来存储一些有层级归属关系的数据。一个最典型的使用树的场景就是浏览器中的 GUI 渲染策略，在你阅读本篇文章的时候，浏览器就已经从 HTML 文件中解析出当前界面上的所有元素，并按照从属关系依次渲染出来，你只需要按下键盘上的 F12 键，就可以看到页面中所有元素的从属关系。

然而，一般意义上的树并不是高度结构化的数据，虽然一颗树的所有信息都可以由根节点遍历出来，但是在没有提前建立索引的情况下，你很难轻易地直接获取树中任意节点的数据。所以人们为了更高的随机索引速度，会更倾向于使用结构稳定的二叉树，尤其是完全二叉树，这是由于完全二叉树的左右子树高度差不超过 1，其每层的节点数量都是 $2^n$ ，所以完全二叉树可以很轻松地存储在线性数组里，然后通过 child = tree[parent_index * 2] 的方式去递推任意孩子节点的索引值。

二叉树的这种特性使得树状结构可以很方便地存储在线性数组中，而本文就将阐述那些和线性数组关系紧密的树状数据结构们。

# 二叉堆

在使用二叉堆之前，你需要知道什么是堆，以及为什么要有堆。

根据维基百科 (opens new window)的描述，堆是一种具有特殊顺序的树，也就是说，堆就是一种树，就像二叉搜索树那样，堆的节点间也有一些大小关系，最常用的堆是最大堆和最小堆，又称大（小）顶堆、大（小）根堆等，以最大堆为例，最大堆中的任意一个节点都比它的子节点的值更大，也就是说，堆的根节点的元素是整个堆的最大值，值得注意的是，这种大小约束只对父节点和子节点生效，而相同层级的节点不需要有大小关系上的约束。

除此之外，堆还必须是一颗完全二叉树，这可以保证对堆的各种操作的均摊复杂度最低。

但是你可能会依然很疑惑，就像我在数据结构课上听到一个全新的数据结构的时候，心里会想：“wow, awesome, 然后呢？”，人们造出来拥有很多复杂结构的事物，并不是完全为了消磨时间，而是因为为了解决问题而不得不让这些工具变得这么复杂，本文中提到的三种数据结构也是如此，它们本身的特性使得它们在解决某些问题时异常好用。

而堆的最常用例子就是优先队列，你可以看到堆的根节点始终是最大值或者最小值，这种最值可以以优先级的形式体现出来，而且这种极值顺序会在动态增减的过程中始终以 $O(log N)$ 的时间复杂度保持着，回想一下，如果你要对一个普通的线性数组取极大值，你就不得不付出 $O(N)$ 的时间复杂度。堆的特性可以大大降低很多需要动态维护优先级的算法的复杂度，比如原始的迪杰斯特拉最短路径算法 (opens new window)的时间复杂度是 $O(N^2)$ ，而使用堆优化后可以达到 $O(N logN)$ 。

现在你知道了什么是堆，以及堆用来解决那些问题，下面就介绍一下堆的常用操作，本文所述的堆均代表二叉堆。

二叉堆通常存储在一个数组中，而堆的插入（insert）、弹出（pop）、取堆顶（top）、删除（remove）操作均可以通过不断交换数组元素来完成，而这些操作均有两个最基础的操作组成：向上调整（up）和向下调整（down），顾名思义，向上调整就是从当前节点开始向根节点遍历，确认该路径上所有节点是否满足“子节点小于（大于）父节点”的顺序，如果不满足，则进行调整；而向下调整则是向叶子节点进行遍历，然后执行同样的操作，写成代码形式（以最大堆为例）：

class Heap:
  def __init__(self):
    self.heap = []

  def swap(self, i, j):
    self.heap[i], self.heap[j] = self.heap[j], self.heap[i]

  def up(self, k):
    '''向上调整第 k 个节点'''
    while k > 0:
        parent_index = (k - 1) // 2
        # 如果子节点大于父节点，则进行调整
        if self.heap[k] > self.heap[parent_index]:
          self.swap(k, parent_index)
        else: break

  def down(self, k)
    '''向下调整第 k 个节点'''
    while k < len(self.heap):
      lchild, rchild = k * 2 + 1, k * 2 + 2
      # 两个子节点取较大节点
      child = lchild if self.heap[lchild] > self.heap[rchild] or rchild > len(self.heap) else rchild
      if child < len(self.heap) and self.heap[child] > self.heap[k]:
        self.swap(child, k)
        k = child
      else: break

可以看到，调整部分的代码的逻辑还是比较清晰的，只需要不断向上或者向下遍历，然后调整对应的值即可。有了两个基础操作，我们可以很方便地完成插入、弹出、取顶、删除等操作了，直接看代码：

class Heap:
  '''省略 up, down 部分的代码'''
  def insert(self, val):
    '''将 val 插入堆中'''
    self.heap.append(val) # 插到数组尾部
    self.up(len(self.heap) - 1) # 然后不断向上调整即可

  def top(self):
    '''获取堆顶值'''
    return None if len(self.heap) == 0 else self.heap[0] # 只需返回数组中的第一个元素即可

  def remove(self, k):
    '''移除第 k 个节点'''
    self.swap(k, len(self.heap) - 1) # 将第 k 个元素交换到尾部
    ret = self.heap.pop() # 将其从数组中删除
    self.up(k) # 确认是否需要向上调整
    self.down(k) # 确认是否需要向下调整
    return ret

  def pop(self):
    '''弹出堆顶'''
    return self.remove(0)

掌握了二叉堆，我们可以很轻松地解决 TOP-k 问题，以及所有需要使用到优先队列的问题，不过值得一提的是，二叉堆作为一种非常常用的数据结构，已经被内置到很多语言的官方库中，在实际使用或者写算法题时，可以直接调包使用，比如 Python3 的 queue.PriorityQueue 和 heapq，以及 C++ 的 priority_queue 等。

# 线段树和树状数组

线段树和树状数组非常像，可以说树状数组就是线段树的一种简化形式，在应对单点修改的区间问题时，树状数组更为简洁好用，但由于使用了 low-bit 技巧，相对来说并没有线段树容易理解，所以本文就先从线段树的原理讲起，再逐步扩展到树状数组。

首先，我们用一个例题来作为切入概念：

来源：303. 区域和检索 - 数组不可变 (opens new window) 给定一个整数数组 nums，求出数组从索引 i 到 j (i ≤ j) 范围内元素的总和，包含 i, j 两点。

下面是一些示例，给定 nums = [-2, 0, 3, -5, 2, -1]，那么对任意的 (i, j) 进行求和：

sumRange(0, 2) -> 1
sumRange(2, 5) -> -1
sumRange(0, 5) -> -3

我们可以很容易地想到 $o(N)$ 的解法，直接对区间 $[i ... j]$ 的数进行求和即可，但是当查询次数很大时，很容易就会超时，而如果我们使用前缀和数组，就可以轻松地在 $o(1)$ 时间内完成任何查询操作，所谓的前缀和数组，就是把前 $i$ 位的数字加起来作为第 $i$ 位的元素值，即 $s[i] = \sum_{j=0}^i a_i$ ，代码表示如下：

def solve(nums):
  s = [0] + nums
  for i in range(1, len(nums) + 1):
    s[i] += s[i - 1]

对于本例，可以算出前缀和数组，为了方便计算，我们往数组前部插入一个零，有了前缀数组，我们就可以使用 $sum(i, j) = s[j + 1] - s[i]$ 来计算任意区间的和了：

s = [0, -2, -2, 1, -4, -2, -3]
sumRange(0, 2) = s[3] - s[0] = 1 - 0 = 1
sumRange(2, 5) = s[6] - s[2] = -3 - (-2) = -1

可以看到，数组的区间信息可以压缩成更紧凑的形式。对于此例题中的静态数组，前缀和应对起来绰绰有余，但是如果在查询过程中数组的数据发生了变化，我们就不得不在每次变化的时候花上 $o(N)$ 的时间开销去更新前缀数组，比如下面这道例题。

来源：307. 区域和检索 - 数组可修改 (opens new window) 给定一个整数数组 nums，求出数组从索引 i 到 j (i ≤ j) 范围内元素的总和，包含 i, j 两点。 update(i, val) 函数可以通过将下标为 i 的数值更新为 val，从而对数列进行修改。

前缀和的本质是维护区间信息，但在应对可变数组时的灵活性太差，所以我们使用更强大的线段树来解决这个问题。

上图表示了线段树的结构，线段树本质上是一个二叉树，它通过不断地把数组二分的方式来从根节点向叶子节点扩展，每个节点都维护了一个区间内的数组信息，这个数组信息可以是区间内的数组和以及最大最小值。

class Node:
    def __init__(self, l, r):
        self.l = l
        self.r = r
        self.lchild = None
        self.rchild = None
        self.val = 0

class NumArray:

    def __init__(self, nums: List[int]):
        self.nums = nums
        self.tree = self.build(0, len(nums) - 1)
 
    def build(self, l, r):
        if l > r: return None
        node = Node(l, r)
        if l == r:
            node.val = self.nums[l]
            return node
        m = (l + r) >> 1
        node.lchild = self.build(l, m)
        node.rchild = self.build(m + 1, r)
        node.val = node.lchild.val + node.rchild.val
        return node

    def update(self, i: int, val: int) -> None:
        d = val - self.nums[i]
        self.nums[i] = val
        q = [self.tree]
        while q:
            node = q.pop()
            if i >= node.l and i <= node.r: node.val += d
            else: continue
            for child in [node.lchild, node.rchild]:
                if child: q.append(child)

    def sumRange(self, i: int, j: int) -> int:
        ret = 0
        q = [self.tree]
        while q:
            node = q.pop()
            if not node: continue
            if node.l >= i and node.r <= j:
                ret += node.val
            elif node.l < j:
                q.append(node.lchild)
            elif node.r > i:
                q.append(node.rchild)
        return ret

[未完待续]