编程实践｜用 MoonBit 实现线段树(三)

引言

在上一篇文章当中我们讨论了如何实现一棵支持区间查询、区间加法的 Immutable 线段树，并且使用了很多 MoonBit 语言当中的独特语法。

而作为"在 MoonBit 实现线段树"系列文章的最后一节，让我们来探讨一下如何实现一个同时支持区间乘法和区间加法的线段树，并且探索 Immutable 的函数式线段树在某些场景中的应用。

区间乘法

基于上一篇文章的线段树，如果我们在当时的需求基础上再加一个需求：实现对 [l, r] 区间内的元素乘以 v，我们应该怎么做呢？

在此之前我们已经学习了 LazyTag，易知其实新加入的操作只影响 Tag 的构成、Tag 与 Tag 的合并、Tag 与 Node 的合并三个部分，让我们来分别思考他们。

Tag 的构成

由于新加入了一个"乘法"的操作，我们将原本 Tag 的结构从 Tag(Int) 拓展为 Tag(Int, Int)，并且使用 label 特性标明乘法 Tag 与加法 Tag，使其在使用时更加轻松。

moonbit 复制代码

enum LazyTag {
  Nil
  Tag(add~ : Int, mul~ : Int)
} derive(Show)

Tag 与 Tag 的合并

原本只有加法的 Tag 与 Tag 的合并是这样的：

moonbit 复制代码

fn op_add(self : LazyTag, v : LazyTag) -> LazyTag {
  match (self, v) {
    (Tag(a), Tag(b)) => Tag(a + b)
    (Nil, t) | (t, Nil) => t
  }
}

我们实际上利用了数学当中的"加法交换律"，即在操作只有加法的情况下，两个 Tag 的合并顺序和最终结果是没有关系的，因为他们所代表的操作均为加法，Tag 合并的结果也只是把他们对节点加上的值相加起来而已。

而一旦引入了"乘法操作"，根据数学定律，乘法的运算顺序优于加法，则其实 Tag 之间的合并次序会出现"先后"，因为乘法总是要比加法先一步计算，所以我们假设 op_add 函数的第一个参数为较老的 Tag，第二个参数为较新的 Tag。则其实我们注意到应该先将新 Tag 的乘法部分 apply 到老 Tag 的乘法/加法部分，再将新 Tag 的加法部分 apply 到老 Tag 的加法部分，代码如下：

moonbit 复制代码

fn op_add(self : LazyTag, v : LazyTag) -> LazyTag {
  match (self, v) {
    (Tag(add=adda, mul=mula), Tag(add=addb, mul=mulb)) =>
      Tag(mul=mula * mulb, add=adda * mulb + addb)
    (Nil, t) | (t, Nil) => t
  }
}

Node 与 Tag 的合并

按照上一节的总结，我们只需要根据"先计算乘法再计算加法"的顺序来把 Tag 的信息转移到 Node 上即可：即先对 sum 的值计算乘法再对其计算加法：

moonbit 复制代码

fn apply(self : Node, v : LazyTag) -> Node {
  match (self, v) {
    (
      Node(data=Data(sum=a, len~), tag~, left~, right~),
      Tag(add~, mul~) as new_tag,
    ) =>
      Node(
        data=Data(sum=a * mul + add * len, ~len),
        tag=tag + new_tag,
        left~,
        right~,
      )
    (_, Nil) => self
    (Nil, _) => Nil
  }
}

这样我们就完成了一棵支持"区间乘法"的线段树，再对其补全一下对应的 API（实际上就是对不同 Tag 修改方式的缩写），需要注意的一点就是在进行 add 的操作时，乘法 Tag 的值应设置为 1，以代表这个 Tag 的乘法部分不影响整体结果：

moonbit 复制代码

fn mul(
  self : Node,
  l : Int,
  r : Int,
  modify_l : Int,
  modify_r : Int,
  value : Int
) -> Node {
  modify(self, l, r, modify_l, modify_r, Tag(add=0, mul=value))
}

fn add(
  self : Node,
  l : Int,
  r : Int,
  modify_l : Int,
  modify_r : Int,
  value : Int
) -> Node {
  modify(self, l, r, modify_l, modify_r, Tag(add=value, mul=1))
}

Immutable 线段树的一些实际应用

上篇文章当中我们介绍过 MoonBit 当中的垃圾回收（GC）机制，而基于这样的内存回收机制，不需要手动精细地管理内存，Immutable 的线段树就可以做到下方这样的内存分布：即每次修改产生新版本时只新建到根的一条路径，而原本基底不变，这就可以支撑我们解决更多有意思的问题。

比如下方这个问题：

给定一个由 n 个整数构成的序列 a，对于指定的闭区间 [l, r] 查询其区间内的第 k 小值。

一个非常简单的想法肯定是先排序之后取出第 k 小值即可，但这样每次处理的复杂度是 O(N Log N) 的（假设使用的排序算法为快速排序）。

而在这里我们可以引入一种叫做可持久化权值线段树的数据结构来更优雅地解决这个问题：

以 Immutable 的方法建立一棵权值线段树（即覆盖整个值域，可以统计某个数字出现的次数的线段树）
接下来遍历这个数列，向该权值线段树的对应位置增加统计值并保存该版本。如遍历到了值 5，则线段树上 5 的位置就应该加一，并且我们额外保存这个新成立的版本。
此时我们容易发现，如果我们希望得到 [l, r] 中多少数字的统计信息，只需要先取出第 r 个版本的线段树进行统计，再减去第 l-1 个版本的线段树统计结果即可，而对于该结果，我们只需要在区间上作"选择左边或者右边"的分治就可以找到题设需求的"第k小值"。

更新操作

我们接下来根据第一篇文章中的"支持"求和的线段树进行更改，将其添加一个单点修改（为某个位置的值+1）操作：

moonbit 复制代码

fn update(self : Node, l : Int, r : Int, k : Int) -> Node {
  match self {
    Node(sum~, left~, right~) => {
      let mid = (l + r) >> 1
      if k <= mid {
        Node(sum=sum + 1, left=update(left, l, mid, k), right~)
      } else {
        Node(sum=sum + 1, left~, right=update(right, mid + 1, r, k))
      }
    }
    Nil => Nil
  }
}

一些工具函数

为了对接下来的流程进行简写，我们暂时定义了一些工具函数来解决需求：

moonbit 复制代码

fn unwrap(self : Node) -> Int {
  match self {
    Node(sum~, ..) => sum
    Nil => 0
  }
}

fn left_tree(self : Node) -> Node {
  match self {
    Node(left~, ..) => left
    Nil => Nil
  }
}
  
fn right_tree(self : Node) -> Node {
  match self {
    Node(right~, ..) => right
    Nil => Nil
  }
}

可持久化权值线段树上的查找

这个操作是整个数据结构实现当中最重要的部分，根据上面的论述，对于查询第 k 小值的操作，我们只要选出 l-1 与 r 两个版本，然后再另外维护一个权值上的二分 l=1 r=len，一步一步地在权值线段树上二分查询即可：

moonbit 复制代码

fn search(left : Node, right : Node, l : Int, r : Int, k : Int) -> Int {
  let mid = (l + r) >> 1
  let x = right.left_tree().unwrap() - left.left_tree().unwrap()
  if l == r {
    return l
  }
  if k <= x {
    search(left.left_tree(), right.left_tree(), l, mid, k)
  } else {
    search(left.right_tree(), right.right_tree(), mid + 1, r, k - x)
  }
}

fn find(versions : Array[Node], l : Int, r : Int, k : Int) -> Int {
  search(versions[l - 1], versions[r], 1, 5, k)
}

fn main {
  // 假设值域为 [1, 5]
  let trees : Array[Node] = []
  let mut tree = build([0, 0, 0, 0, 0][:])
  trees.push(tree)
  let arr = [1, 2, 3, 4, 5, 5, 4, 3, 2, 1]
  for x in arr {
    tree = tree.update(1, 5, x)
    trees.push(tree)
  }
  println(find(trees, 1, 7, 5))
}

总结

在这篇文章中我们首先完成了 Tag 结构更加复杂的、支持区间乘法的线段树，然后又探索了 Immutable 线段树的一个优秀的应用场景------使用可持久化权值线段树求静态区间最小值，并在此过程中充分地利用了 MoonBit 的 GC 优势与丰富的语言特性。

自此"用 MoonBit 实现线段树"系列文章完结，感兴趣的读者可以自行继续阅读并研究线段树的更高级知识，如矩阵乘法线段树、线段树分裂等。

可持久化线段树的例题与内存结构参考自 oiwiki.com/ds/persiste...