Code. Simple segment tree implementation in MoonBit [blog/fenwick/segtree_code]

下面代码展示了一个在 MoonBit 中实现的简单线段树,其中使用了一些处理索引范围的工具函数,如图 8 所示。我们将线段树存储为一个递归的代数数据类型。 并使用与前一节中递归描述直接对应的代码实现了 updaterq;随后,getset 亦可基于它们来实现。将此代码推广至适用于存储来自任意交换幺半群(若不需要 set 操作)或任意阿贝尔群(即带逆元的交换幺半群,若需要 set 操作)的值的线段树亦非难事——但为简单起见,我们在此保持原状,因为这种推广对我们的主线故事并无增益。

enum SegTree {
  Empty
  Branch(Int, Range, SegTree, SegTree)
} derive(Eq, Show)

pub fn SegTree::update(self : SegTree, i : Index, v : Int) -> SegTree {
  match self {
    Empty => Empty
    Branch(a, rng, l, r) if rng.contains(i) =>
      Branch(a + v, rng, l.update(i, v), r.update(i, v))
    _ => self
  }
}

fn rq(self : SegTree, q : Range) -> Int {
  match self {
    Empty => 0
    Branch(a, rng, l, r) =>
      if disjoint(rng, q) {
        0
      } else if subset(rng, q) {
        a
      } else {
        l.rq(q) + r.rq(q)
      }
  }
}

pub fn SegTree::get(self : SegTree, i : Index) -> Int {
  self.rq((i, i))
}

pub fn SegTree::set(self : SegTree, i : Index, v : Int) -> SegTree {
  self.update(i, v - self.get(i))
}

区间工具函数:

typealias Index = Int

priv type Range (Index, Index) derive(Eq, Show)

fn subset(r1 : Range, r2 : Range) -> Bool {
  let (a1, b1) = r1._
  let (a2, b2) = r2._
  a1 >= a2 && b1 <= b2
}

fn contains(self : Range, i : Index) -> Bool {
  subset((i, i), self)
}

fn disjoint(r1 : Range, r2 : Range) -> Bool {
  let (a1, b1) = r1._
  let (a2, b2) = r2._
  b1 < a2 || b2 < a1
}