线段树-树状树学习及应用

Narcissus 收录于算法学习

2022-09-07 约 4212 字预计阅读 9 分钟

假设我们有一个数组：

数组不变，求区间和：「前缀和」、「树状数组」、「线段树」
多次修改某个数（单点），求区间和：「树状数组」、「线段树」
多次修改某个区间，输出最终结果：「差分」
多次修改某个区间，求区间和：「线段树」、「树状数组」（看修改区间范围大小）
多次将某个区间变成同一个数，求区间和：「线段树」、「树状数组」（看修改区间范围大小）

但是因为线段树代码量很大，所以我们应该按如下优先级进行考虑：

简单求区间和，用「前缀和」;

多次将某个区间变成同一个数，用「线段树」;

其他情况，用「树状数组」

注意：树状数组仅仅适合于求区间和的问题，而线段树适合于区间问题。

例题：LeetCode699:本题就是将区间变成同一个数，求最大值，注意并不是相加。

差分数组

前缀和主要适用的场景是原始数组不会被修改的情况下，频繁查询某个区间的累加和。

差分数组主要适用的场景是多次修改某个区间，输出最终结果。

构造差分数组，diff[i] = nums[i] - nums[i-1]：

如果需要对区间[i - j]的元素全部加上k，则只需要diff[i] += k diff[j+1] -= k即可。

通过差分数组还原原数组，可以看出原理，diff[i] += k 相当于给nums[i…]所有元素都加上k。

代码如下：

n := len(nums)
diff := make([]int, n)

// 构建差分数组
diff[0] = nums[0]
for i := 1; i < n; i++ {
  diff[i] = nums[i] - nums[i-1]
}

// 根据差分数组还原原数组
nums[0] = diff[0]
for i := 1; i < n; i++ {
  nums[i] = nums[i-1] + diff[i]
}

LeetCode题目1109

树状数组

1. 介绍

树状数组就是用数组来模拟树形结构。可以用来解决大部分基于区间上的更新以及求和问题。树状数组可以解决的问题都可以用线段树解决。但是线段树代码很长，常数很大，所以优先考虑树状树。修改和查询的复杂度都是O(logN)，树状数组本质就是二进制规律的应用。

树形结构如下：

黑色数组代表原来的数组（下面用A[i]代替），红色结构代表我们的树状数组(下面用C[i]代替)，则有如下：

C[1] = A[1]

C[2] = A[1] + A[2]

C[3] = A[3]

C[4] = A[1] + A[2] + A[3] + A[4]

C[5] = A[5]

C[6] = A[5] + A[6]

C[7] = A[7]

C[8] = A[1] + A[2] + A[3] + A[4] + A[5] + A[6] + A[7] + A[8]

可以发现规律为：C[i] = A[i-2^k+1]+A[i-2^k+2]+…+A[k]; (k为i的二进制中从最低位到高位连续零的长度)

例如，i=8（1000）的时候，k=3。

如何求和？例如找前7项和，那么SUM=C[7]+C[6]+C[4];

则可以发现求和规律：SUMi = C[i]+C[i-2^k1]+C[(i-2^k1)-2^k2]+……

注意：2^k怎么算呢？前辈的智慧：i&(-i)，并且有一个专门的名称，lowbit。

2. 建立树状数组

如果我们更新A[i]，则会影响到包含A[i]的所有位置，A[i]包含于C[i+2^k]，C[(i+2^k)+2^k]……

单点更新，区间查询

代码如下：

	tree := make([]int, n+1) // int64

	// 更新，在i位置加上val
	add := func(i int, val int) {
		for ; i < len(tree); i += i & -i {
			tree[i] += val
		}
	}

	// 求原数组[1 - i]的和
	sum := func(i int) (res int) {
		for ; i > 0; i &= i - 1 {
			res += tree[i]
		}
		return
	}

	// 求原数组[l - r]的和
	query := func(l, r int) int {
		return sum(r) - sum(l-1)
	}

结构体写法

// 结构体写法
type fenwick struct {
	tree []int64
}

func newFenwickTree(n int) fenwick {
	return fenwick{make([]int64, n+1)}
}

// 位置 i 增加 val
// 1<=i<=n
func (f fenwick) add(i int, val int64) {
	for ; i < len(f.tree); i += i & -i {
		f.tree[i] += val
	}
}

// 求前缀和 [0,i]
// 0<=i<=n
func (f fenwick) sum(i int) (res int64) {
	for ; i > 0; i &= i - 1 {
		res += f.tree[i]
	}
	return
}

// 求区间和 [l,r]
// 1<=l<=r<=n
func (f fenwick) query(l, r int) int64 {
	return f.sum(r) - f.sum(l-1)
}

LeetCode题目：307

区间更新，单点查询

如果需要将[x - y]区间内所有值均加上k，如果采用上述方式，复杂度肯定不行。这个时候就不需要用原数据建树，而是用差分数组建树。

则当对[x - y]区间内所有值做修改，只需要修改差分数组D[x]、D[y+1]即可。

此时求差分数组D[1 - i]的和就是求元素组A[i]。D[i] + D[i-1] +…+D[0] = A[i] - A[i-1] + A[i-1] - A[i-2] +…+A[1] - A[0] = A[i] - A[0] = A[i]

注意：为了构建差分数组，规定A[0] = 0

代码如下：

tree := make([]int, n+1) // int64

// 更新，在i位置加上val
add := func(i int, val int) {
  for ; i < len(tree); i += i & -i {
    tree[i] += val
  }
}

// 更新区间[l, r]均加上val
// 注意：r+1即使超过n也没关系，因为不会用到
addRange := func(l, r int, val int) {
  add(l, val)
  add(r+1, -val)
}

// 求差分数组[1 - i]的和，即求A[i]
sum := func(i int) (res int) {
  for ; i > 0; i &= i - 1 {
    res += tree[i]
  }
  return
}

区间更新，区间查询

因为仍然需要区间更新，所以依然使用差分数组建树。推到如下：

代码如下：

tree1 := make([]int, n+1) // D[i]
tree2 := make([]int, n+1) // D[i]*(i-1)

// 更新，在i位置加上val
add := func(i int, val int) {
  x := i - 1
  for ; i < len(tree); i += i & -i {
    tree1[i] += val
    tree2[i] += val*x
  }
}

// 更新区间[l, r]均加上val
// 注意：r+1即使超过n也没关系，因为不会用到
addRange := func(l, r int, val int) {
  add(l, val)
  add(r+1, -val)
}

// 求原数组[1 - i]的和
sum := func(i int) (res int) {
  x := i
  for ; i > 0; i &= i - 1 {
    res += x*tree1[i] - tree2[i]
  }
  return
}

// 求原数组区间[l - r]的和
query := func(l, r int) int {
  return sum(r) - sum(l-1)
}

线段树

树状数组能解决的问题，线段树都能解决。线段树应用范围很广，主要用于解决区间问题，但远不止如此。

线段树是一种二叉搜索树，可以在线维护修改以及查询区间上的最值、求和。

线段树基本内容

首先需要明白：每个节点存什么？节点下标是什么？如何建树？

例如，A[1:6] = {1, 8, 6, 4, 3, 5} 如下区间最大值情况：

可以发现，每个叶子节点的值就是数组的值，每个非叶子节点的度都是二，且左右两个孩子分别存储父亲一半的区间。每个父亲的存储值也就是两个孩子存储的值的最大值。

最重要的问题是如何存储区间？以及如何快速找到没叶子结点的孩子以及非根节点的父亲呢？

对于上述线段树，我们增加绿色数字为每个节点的下标：

为什么叶子下标从9跳到了12？因为中间也有两个空间，这也是为什么无优化的线段树建树需要22^k空间，一本会开到4n的空间防止RE。

节点k的左子树下标为k << 1 右子树的下标为k<<1 | 1

常使用递归的方式建树。

线段树基本操作

1. 单点更新，区间查询

更新和查询的时间复杂度是O(logN)

对于上述线段树，我们把a[3]+7，更新后线段树应该为：

我们发现，无论更新哪一个叶子节点，最终都会到根节点，而把这个过程逆过来就是从根节点开始，找到左子树还是右子树包含需要更新的叶子节点，往下更新即可。所以采用递归的方式实现线段树更新

对于上述线段树，比如现在我们要查询[2, 5]区间的最值，可以看看那些区间是[2, 5]的真子集，

一共有五个区间，而区间[4, 5]已经包含其儿子节点的信息，所以一共需要查询3个区间。我们依然从根开始递归，如果当前节点是需要查询的区间的真子集，则返回节点信息并且停止递归。时间复杂度是O(logN)。

代码如下：

// l 和 r 也可以写到方法参数上，实测二者在执行效率上无异
// 考虑到 debug 和 bug free 上的优点，写到结构体参数中
type seg []struct {
	l, r int
	val  int
}

func newSegmentTree(a []int) seg {
	t := make(seg, 4*len(a))
	t.build(a, 1, 1, len(a))
	return t
}

// 单点更新：build 和 update 通用
func (t seg) set(o, val int) {
	t[o].val = val
}

// 合并两个节点上的数据：maintain 和 query 通用
// 要求操作满足区间可加性
// 例如 + * | & ^ min max gcd mulMatrix 摩尔投票 最大子段和 ...
func (seg) op(a, b int) int {
  // 示例为max操作
	if a > b {
		return a
	}
	return b
}

// 更新节点o
func (t seg) maintain(o int) {
	lo, ro := t[o<<1], t[o<<1|1]
	t[o].val = t.op(lo.val, ro.val)
}

// o为当前需要建立的节点，l,r分别表示当前节点表示区间的左右端点
func (t seg) build(a []int, o, l, r int) {
	t[o].l, t[o].r = l, r
	if l == r {
		t.set(o, a[l-1])
		return
	}
	m := (l + r) >> 1
	t.build(a, o<<1, l, m)	//递归构建左儿子节点
	t.build(a, o<<1|1, m+1, r)
	t.maintain(o) // 更新父节点
}

// o表示当前更新到的结点，从o=1开始更新
// i 表示要更新的节点，val表示要更新的值
func (t seg) update(o, i, val int) {
	if t[o].l == t[o].r {
		t.set(o, val)
		return
	}
	if m := (t[o].l + t[o].r) >> 1; i <= m {
		t.update(o<<1, i, val)
	} else {
		t.update(o<<1|1, i, val)
	}
	t.maintain(o)
}

// o=1 表示从根开始递归
// [l,r] 1<=l<=r<=n
func (t seg) query(o, l, r int) int {
	if l <= t[o].l && t[o].r <= r {
		return t[o].val
	}
	m := (t[o].l + t[o].r) >> 1 // 左儿子区间[t[o].l, m] 右儿子区间[m+1, t[o].r]
	if r <= m {
		return t.query(o<<1, l, r)
	}
	if m < l {
		return t.query(o<<1|1, l, r)
	}
	vl := t.query(o<<1, l, r)
	vr := t.query(o<<1|1, l, r)
	return t.op(vl, vr)
}

// 查询整个区间的信息
func (t seg) queryAll() int { 
  return t[1].val 
}

2. 区间更新区间查询

树状数组中区间更新我们采用了差分的思想，而线段树中，我们将引入lazy_tag（懒惰标记），线段树进行区间更新的时候，每次更新只更新到更新区间完全覆盖线段树节点区间为止，这样会导致被更新节点的子孙节点的区间得不到更新信息，所以在被更新节点上打上一个标记。等下次访问到这个节点的子节点时再讲这个标记传递给子节点，所以也可以叫延迟标记。

也就是说，在递归更新过程中，更新到节点区间为需要更新区间的真子集时不再递归更新，下次如果需要用到再更新，所以时间复杂度也是O(logN)。

type lazySeg []struct {
	l, r int
	todo int64
	sum  int64
}

// a 从 0 开始
func newLazySegmentTree(a []int64) lazySeg {
	t := make(lazySeg, 4*len(a))
	t.build(a, 1, 1, len(a))
	return t
}

func (lazySeg) op(a, b int64) int64 {
	return a + b // % mod
}

func (t lazySeg) maintain(o int) {
	lo, ro := t[o<<1], t[o<<1|1]
	t[o].sum = t.op(lo.sum, ro.sum)
}

func (t lazySeg) build(a []int64, o, l, r int) {
	t[o].l, t[o].r = l, r
	if l == r {
		t[o].sum = a[l-1]
		return
	}
	m := (l + r) >> 1
	t.build(a, o<<1, l, m)
	t.build(a, o<<1|1, m+1, r)
	t.maintain(o)
}

func (t lazySeg) do(o int, add int64) {
	to := &t[o]
	to.todo += add                     // % mod
	to.sum += int64(to.r-to.l+1) * add // % mod
}

func (t lazySeg) spread(o int) {
	if add := t[o].todo; add != 0 {
		t.do(o<<1, add)
		t.do(o<<1|1, add)
		t[o].todo = 0
	}
}

// o=1  [l,r] 1<=l<=r<=n
func (t lazySeg) update(o, l, r int, add int64) {
	if l <= t[o].l && t[o].r <= r {
		t.do(o, add)
		return
	}
	t.spread(o)
	m := (t[o].l + t[o].r) >> 1
	if l <= m {
		t.update(o<<1, l, r, add)
	}
	if m < r {
		t.update(o<<1|1, l, r, add)
	}
	t.maintain(o)
}

// o=1  [l,r] 1<=l<=r<=n
func (t lazySeg) query(o, l, r int) int64 {
	if l <= t[o].l && t[o].r <= r {
		return t[o].sum
	}
	t.spread(o)
	m := (t[o].l + t[o].r) >> 1
	if r <= m {
		return t.query(o<<1, l, r)
	}
	if m < l {
		return t.query(o<<1|1, l, r)
	}
	vl := t.query(o<<1, l, r)
	vr := t.query(o<<1|1, l, r)
	return t.op(vl, vr)
}

3. 动态开点

常规线段树，一般以满二叉树方式存储，浪费了很多空间，有时需要动态开点的方式，动态开点的优势就是按需开点，因此一般采用链表方式实现而不是数组方式实现。

3.1 单点修改的动态开点

代码如下

// 创建根节点
rt := &stNode{l: 1, r: 1e9}

type stNode struct {
	lo, ro *stNode
	l, r   int
	sum    int64
}

func (o *stNode) get() int64 {
	if o != nil {
		return o.sum
	}
	return 0 // inf
}

func (stNode) op(a, b int64) int64 {
	return a + b
}

func (o *stNode) maintain() {
	o.sum = o.op(o.lo.get(), o.ro.get())
}

func (o *stNode) update(i int, add int64) {
	if o.l == o.r {
		o.sum += add
		return
	}
	m := (o.l + o.r) >> 1
	if i <= m {
		if o.lo == nil {
			o.lo = &stNode{l: o.l, r: m}
		}
		o.lo.update(i, add)
	} else {
		if o.ro == nil {
			o.ro = &stNode{l: m + 1, r: o.r}
		}
		o.ro.update(i, add)
	}
	o.maintain()
}

func (o *stNode) query(l, r int) int64 {
	if o == nil || l > o.r || r < o.l {
		return 0 // inf
	}
	if l <= o.l && o.r <= r {
		return o.sum
	}
	return o.op(o.lo.query(l, r), o.ro.query(l, r))
}

3.2 区间修改的动态开点

代码如下

// 创建根节点
rt := &lazyNode{l: 1, r: 1e9}

type lazyNode struct {
	lo, ro *lazyNode
	l, r   int
	sum    int64
	todo   int64
}

func (o *lazyNode) get() int64 {
	if o != nil {
		return o.sum
	}
	return 0 // inf
}

func (lazyNode) op(a, b int64) int64 {
	return a + b
}

func (o *lazyNode) maintain() {
	o.sum = o.op(o.lo.get(), o.ro.get())
}

func (o *lazyNode) do(add int64) {
	o.todo += add                   // % mod
	o.sum += int64(o.r-o.l+1) * add // % mod
}

func (o *lazyNode) spread() {
	m := (o.l + o.r) >> 1
  // 动态开点
	if o.lo == nil {
		o.lo = &lazyNode{l: o.l, r: m}
	}
	if o.ro == nil {
		o.ro = &lazyNode{l: m + 1, r: o.r}
	}
	if add := o.todo; add != 0 {
		o.lo.do(add)
		o.ro.do(add)
		o.todo = 0 // -1
	}
}

func (o *lazyNode) update(l, r int, add int64) {
	if l <= o.l && o.r <= r {
		o.do(add)
		return
	}
	o.spread()
	m := (o.l + o.r) >> 1
	if l <= m {
		o.lo.update(l, r, add)
	}
	if m < r {
		o.ro.update(l, r, add)
	}
	o.maintain()
}

func (o *lazyNode) query(l, r int) int64 {
	// 对于不在线段树中的点，应按照题意来返回
	if o == nil || l > o.r || r < o.l {
		return 0
	}
	if l <= o.l && o.r <= r {
		return o.sum
	}
	o.spread()
	return o.op(o.lo.query(l, r), o.ro.query(l, r))
}

4. 离散化

离散化常用于二维状态在一维线段树建树的技巧，所谓离散化就是将无限的个体映射到有限的个体中，提高算法效率，常用Hash实现。

树状数组有时也需要离散化。例题LeetCode327。

目录

线段树-树状树学习及应用

差分数组

树状数组

1. 介绍

2. 建立树状数组

单点更新，区间查询

区间更新，单点查询

区间更新，区间查询

线段树

线段树基本内容

线段树基本操作

1. 单点更新，区间查询

2. 区间更新区间查询

3. 动态开点

3.1 单点修改的动态开点

3.2 区间修改的动态开点

4. 离散化