Go实现动态开点线段树

1、线段树介绍

线段树是一种用于高效处理区间查询和区间更新的数据结构,当我们需要解决一个频繁更新区间值的问题的时候,就可以采用线段树的结构进行解决。线段树的核心思想是将区间分为多个子区间进行管理,越往下区间范围越小,根节点表示整个线段树能表示的区间。

本文记录使用Go实现动态开点线段树的方式,该模板的线段树用于解决区间求和问题,还有求解区间最小值、最大值的线段树可以进行微调修改即可。

区间查询、区间更新的时间复杂度均为O(logN)

2、动态开点线段树实现

动态开点的核心在于,需要缩小范围,即进入子节点的时候再进行创建,相对于使用数组来实现线段树,可以更大的减小空间开销。

1、线段树节点

一个节点需要记录它的左子节点、右子节点、当前节点表示的区间的和val,以及暂未下推给子节点的懒惰值lazy

type SegTreeNode struct {
	lazy  int
	val   int
	left  *SegTreeNode
	right *SegTreeNode
}

2、线段树的创建

整个线段树只需要记录一个根节点以及该线段树表示的区间上届。

type SegTree struct {
	//线段树的范围,0~N
	N    int
	root *SegTreeNode
}

// 创建线段树
func CreateSegTree(n int) *SegTree {
	return &SegTree{
		N: n,
		root: &SegTreeNode{
			lazy:  0,
			val:   0,
			left:  nil,
			right: nil,
		},
	}
}

3、递归上推

当更新完了子节点后,回到当前节点的时候,需要更新当前节点的值,表示从树的底部上推值。

// 递归上推
func (ST *SegTree) Pushup(node *SegTreeNode) {
	node.val = node.left.val + node.right.val
}

4、懒惰下推

当需要缩小查找区间的时候,需要向下查找,这时候要先把懒惰值下推,防止查找出错误的结果,也防止子节点还未创建。

// 同步下推
func (ST *SegTree) Pushdown(node *SegTreeNode, leftnum, rightnum int) {
	//创建左右节点
	if node.left == nil {
		node.left = new(SegTreeNode)
	}
	if node.right == nil {
		node.right = new(SegTreeNode)
	}
	//下推节点懒惰标记
	if node.lazy == 0 {
		return
	}
	node.left.val += leftnum * node.lazy
	node.right.val += rightnum * node.lazy
	//下推
	node.left.lazy += node.lazy
	node.right.lazy += node.lazy
	//置零
	node.lazy = 0
}

首先先创建左右节点,如果没有需要下推的懒惰标记则直接返回。否则就更新左右节点的vallazy

5、更新操作

// 更新操作,更新[left,right]区间的值,start和end是当前处在区间
func (ST *SegTree) Update(node *SegTreeNode, start, end, left, right, val int) {
	if left <= start && end <= right {
		//锁定区间,进行更新
		node.val += (end - start + 1) * val
		node.lazy += val
		return
	}
	//缩小区间
	mid := (start + end) / 2
	//需要找到子节点,先下推懒惰标记
	ST.Pushdown(node, mid-start+1, end-mid)
	if mid >= left {
		ST.Update(node.left, start, mid, left, right, val)
	}
	if mid+1 <= right {
		ST.Update(node.right, mid+1, end, left, right, val)
	}
	//递归
	ST.Pushup(node)
}

leftright表示要更新的区间,而startend表示当前区间。如果当前区间处在需要更新的区间内,则直接更新区间值以及懒惰值,然后直接返回即可,此时不需要继续更新下面节点的值,这是动态开点的关键所在。

若当前区间并未完全处在需要更新的区间内,则二分该区间,缩小范围进行更新。

例如在一次操作需要更新的是[30,40]范围的值,而当前区间处在[25,50]中,当前区间并未完全处在更新区间,则二分为[25,37][38,50],左区间和右区间均和需要更新的区间存在交集,那么就往下更新,直到更新区间包含当前区间。

在更新完后,进行一次上推。

6、查询操作

与更新操作类似,只需要一个ans来记录答案并且返回。

// 查询操作,返回区间的值
func (ST *SegTree) Query(node *SegTreeNode, start, end, left, right int) int {
	if left <= start && end <= right {
		return node.val
	}
	mid := (start + end) / 2
	ST.Pushdown(node, mid-start+1, end-mid)
	ans := 0
	if left <= mid {
		ans += ST.Query(node.left, start, mid, left, right)
	}
	if mid+1 <= right {
		ans += ST.Query(node.right, mid+1, end, left, right)
	}
	return ans
}

3、尝试题目

LeetCode我的日程表安排I

[LeetCode我的日程表安排III](732. 我的日程安排表 III - 力扣(LeetCode))

2502. 设计内存分配器 - 力扣(LeetCode)

3.1、LC2502设计内存分配器

改题可以用模拟的做法,但是如果数据量再大一点,用线段树可以达到更优的解。

本题可以用线段树记录区间连续0字段的最大个数,用于完成对内存的分配,我们先看需要什么信息:

  • 想要分配一段内存,我们需要知道满足size个连续0的线段树的起始位置
  • 想要释放一段内存,直接对这段内存的位置置为0

位置我们可以用区间不断缩小来找到,而想要知道该区间最大的连续0个数是多少,我们需要维护一个区间的max0变量。

考虑max0如何更新,受到以下状态影响:

  • 左节点区间的max0
  • 右节点区间的max0
  • 若左节点的末尾是0,那么可以由左节点的末尾连续0拼接右节点头部的连续0。

max0的更新取上述三值的最大值。

从第三个状态可以知道,我们还需要维护一个节点的头部连续0个数pre0,以及尾部连续0个数suf0

再来看如何找到满足size的区间起点:

  • 因为要求左往右找,因此需要先遍历左区间,若有符合要求的区间,则返回区间左端点
  • 否则,检查左区间的suf0+右区间的pre0是否满足
  • 递归右区间

递归的边界:

  • 若当前区间的max0 < size,那么该区间可以直接返回,因为不满足要求
  • l==r,则直接返回左端点

代码如下,这里将原本上面模板的update拆分成modifyupdate

type SegNode struct {
	suf0, pre0, max0 int
	lazy             int
	left, right      *SegNode
	all0             bool
}
type SegTree struct {
	N    int
	root *SegNode
}

func CreateSegNode(n int) *SegNode {
	return &SegNode{
		suf0: n,
		pre0: n,
		max0: n,
		all0: true,
		lazy: -1,
	}
}
func CreateSegTree(n int) *SegTree {
	return &SegTree{
		N: n,
		root: &SegNode{
			suf0: n+1,
			pre0: n+1,
			max0: n+1,
			all0: true,
			lazy: -1,
		},
	}
}

// 上推
func PushUp(node *SegNode) {
	node.pre0 = node.left.pre0
	node.suf0 = node.right.suf0
	if node.left.all0 {
		node.pre0 += node.right.pre0
	}
	if node.right.all0 {
		node.suf0 += node.left.suf0
	}
	node.max0 = max(node.left.max0, node.right.max0, node.left.suf0+node.right.pre0)
	node.all0 = node.left.all0 && node.right.all0
}
func PushDown(node *SegNode, l, r int) {
	mid := (l + r) / 2
	//lazy标记为1表示所有置1,为0表示置0,为-1表示没操作
	if node.left == nil {
		node.left = CreateSegNode(mid - l + 1)
	}
	if node.right == nil {
		node.right = CreateSegNode(r - mid)
	}
	if node.lazy != -1 {
		Update(node.left, l, mid, node.lazy)
		Update(node.right, mid+1, r, node.lazy)
		//下推完成
		node.lazy = -1
	}
}
func Update(node *SegNode, l, r, val int) {
	//当前区间为[l,r],直接更新值,val为0或1
	if val == 1 {
		node.all0 = false
		node.pre0 = 0
		node.suf0 = 0
		node.all0 = false
		node.max0 = 0
	} else {
		node.all0 = true
		node.pre0 = r - l + 1
		node.suf0 = node.pre0
		node.max0 = node.pre0
	}
	node.lazy = val
}

// 将[l,r]赋值为val,当前为[curl,curr]
func Modify(node *SegNode, curl, curr, l, r int, val int) {
	if l <= curl && curr <= r {
		Update(node, curl, curr, val)
		return
	}
	//下推
	PushDown(node, curl, curr)
	mid := (curl + curr) / 2
	if mid >= l {
		Modify(node.left, curl, mid, l, r, val)
	}
	if mid < r {
		Modify(node.right, mid+1, curr, l, r, val)
	}
	//上推
	PushUp(node)
}

// 查询第一个连续0的起始点
// 将线段树二分,先找node的左区间,找到返回下标,否则返回-1
func FindFirstPos(node *SegNode, l, r int, len int) int {
	if node.max0 < len {
		return -1
	}
	if l == r {
		return l
	}
	mid := (l + r) / 2
	//下推
	PushDown(node, l, r)
	idx := FindFirstPos(node.left, l, mid, len)
	//ids必定为-1,左区间刚好不能满足,查看能否左区间末尾和右区间前端连续的0满足,可以则返回,否则递归右子树
	if idx < 0 {
		if node.left.suf0+node.right.pre0 >= len {
			return mid - node.left.suf0 + 1
		}
		idx = FindFirstPos(node.right, mid+1, r, len)
	}
	return idx

}

type interval struct {
	l, r int
}
type Allocator struct {
	record  map[int][]interval
	SegTree *SegTree
}

func Constructor(n int) Allocator {
	return Allocator{
		record:  make(map[int][]interval),
		SegTree: CreateSegTree(n-1),
	}
}

func (this *Allocator) Allocate(size int, mID int) int {
	idx := FindFirstPos(this.SegTree.root, 0, this.SegTree.N, size)
	if idx == -1 {
		return -1
	}
	//记录起始区间
	this.record[mID] = append(this.record[mID], interval{idx, idx + size - 1})
	//分配内存
	Modify(this.SegTree.root, 0, this.SegTree.N, idx, idx+size-1, 1)
	return idx
}

func (this *Allocator) FreeMemory(mID int) int {
	ans := 0
	for _, v := range this.record[mID] {
		ans += v.r - v.l + 1
		Modify(this.SegTree.root, 0, this.SegTree.N, v.l, v.r, 0)
	}
	delete(this.record, mID)
	return ans
}

3.2、LC2313由单个字符重复的最长子字符串

思路和上题大体一致,要单独记录区间最前面和最后面的字符。

//线段树:
type SegNode struct {
	Last, First      int //首位字符
	Prec, Sufc, Maxc int //头连续字符个数、尾连续字符个数、最大连续字符个数
	Lazy             int //将区间置为lazy值,-1无动作
	Left, Right      *SegNode
	Allc             bool //标识是否区间全为相同的字母
}
type SegTree struct {
	N    int //表示涵盖区间[0,N),有N个节点
	Root *SegNode
}

func CreateSegNode(N int) *SegNode {
	return &SegNode{
		Last:  -1,
		First: -1,
		Prec:  N,
		Sufc:  N,
		Maxc:  N,
		Lazy:  -1,
		Allc:  true,
	}
}
func CreateSegTree(N int) *SegTree {
	return &SegTree{
		N:    N,
		Root: CreateSegNode(N),
	}
}

//上推
func PushUp(node *SegNode) {
	node.First = node.Left.First
	node.Last = node.Right.Last
	node.Prec = node.Left.Prec
	node.Sufc = node.Right.Sufc
	if node.Left.Allc && node.Left.Last == node.Right.First {
		node.Prec += node.Right.Prec
	}
	if node.Right.Allc && node.Left.Last == node.Right.First {
		node.Sufc += node.Left.Sufc
	}
	node.Maxc = max(node.Left.Maxc, node.Right.Maxc)
	if node.Left.Last == node.Right.First {
		node.Maxc = max(node.Maxc, node.Left.Sufc+node.Right.Prec)
	}
	node.Allc = node.Left.Allc && node.Right.Allc && node.Left.First == node.Right.First
}

//下推
func PushDown(node *SegNode, l, r int) {
	mid := (l + r) / 2
	if node.Left == nil {
		node.Left = CreateSegNode(mid - l + 1)
	}
	if node.Right == nil {
		node.Right = CreateSegNode(r - mid)
	}
	if node.Lazy != -1 {
		Update(node.Left, l, mid, node.Lazy)
		Update(node.Right, mid+1, r, node.Lazy)
		node.Lazy = -1
	}
}

//更新
func Update(node *SegNode, l, r, val int) {
	//更新[l,r]为val
	node.Lazy = val
	node.Allc = true
	node.First = val
	node.Last = val
	node.Prec = r - l + 1
	node.Sufc = r - l + 1
	node.Maxc = r - l + 1
}

//更新区间为val,当前位于cr,cl
func Modify(node *SegNode, l, r, cl, cr, val int) {
	if l <= cl && cr <= r {
		Update(node, cl, cr, val)
		return
	}
	//缩小区间
	PushDown(node, cl, cr)
	mid := (cl + cr) / 2
	if mid >= l {
		Modify(node.Left, l, r, cl, mid, val)
	}
	if mid < r {
		Modify(node.Right, l, r, mid+1, cr, val)
	}
	PushUp(node)
}
func longestRepeating(s string, queryCharacters string, queryIndices []int) []int {
	t := CreateSegTree(len(s))
	ans := []int{}
	for i := 0; i < len(s); i++ {
		cnt := 1
		for i < len(s)-1 && s[i] == s[i+1] {
			i++
			cnt++
		}
		//初始化
		Modify(t.Root, i-cnt+1, i, 0, len(s)-1, int(s[i]-'a'))
	}
	for i := range queryCharacters {
		Modify(t.Root, queryIndices[i], queryIndices[i], 0, len(s)-1, int(queryCharacters[i]-'a'))
		ans = append(ans, t.Root.Maxc)
	}
	return ans
}
posted @ 2025-02-25 20:04  MelonTe  阅读(196)  评论(0)    收藏  举报