# 树和数组的千层套路

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

目录:

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

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

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

# 二叉堆

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

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

image

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

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

而堆的最常用例子就是优先队列,你可以看到堆的根节点始终是最大值或者最小值,这种最值可以以优先级的形式体现出来,而且这种极值顺序会在动态增减的过程中始终以 O(logN)O(log N) 的时间复杂度保持着,回想一下,如果你要对一个普通的线性数组取极大值,你就不得不付出 O(N)O(N) 的时间复杂度。堆的特性可以大大降低很多需要动态维护优先级的算法的复杂度,比如原始的迪杰斯特拉最短路径算法 (opens new window)的时间复杂度是 O(N2)O(N^2),而使用堆优化后可以达到 O(NlogN)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.PriorityQueueheapq,以及 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)o(N) 的解法,直接对区间 [i...j][i ... j] 的数进行求和即可,但是当查询次数很大时,很容易就会超时,而如果我们使用前缀和数组,就可以轻松地在 o(1)o(1) 时间内完成任何查询操作,所谓的前缀和数组,就是把前 ii 位的数字加起来作为第 ii 位的元素值,即 s[i]=j=0iais[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]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)o(N) 的 时间开销去更新前缀数组,比如下面这道例题。

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

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

image

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

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

[未完待续]