静谧的夜最适合刷算法题了。刷着刷着发现了一个好玩的数据结构叫做线段树,据说是算法竞赛的常客哦,于是就自己写一个简单的玩玩。
原理
想了解线段树的同志们请移步,里边有原理和示意图:
百度百科,线段树
需求和痛点
我用它来主要是为了快速找到数组某区间内的数字和,并且在修改数组某几个元素之后再次找区间内的数字和。可想而知,我有两个需求:求和,修改。
正常情况下,因为我们并不想给数组排序,那么,我们可以用O(1)的复杂度进行修改,用O(n)的复杂度遍历区间来实现求和。但是,当我们需要频繁进行求和的操作时,看起来很美的O(n)就变成了墙上的蚊子血,再也不是当初的红玫瑰了,于是,我们只能优化它。
我选择了线段树,典型的空间换时间。我用数组实现了二叉的线段树,用了4倍的额外空间。这颗二叉树的叶子节点是原始数组的各个元素,非叶子节点存储的为某段区间的和,一直到根节点逐步合并区间,根节点正好就是原始数组从头到尾的最大区间的元素和,具体的可以看看代码,个人觉得代码比文字更直观哈哈。这样的处理将求和操作变成了O(logN),而相应的修改操作也增加到了O(logN),两个对数级别往往总是要好过一个常数的和一个线性的,不是么?
Talk is cheap. Show me the code.
首先是构建线段树,明白原理之后就像二叉树一样直接递归着搞就好啦:
type segmentTree struct {
data []int //原始数组
tree []int //线段树数组
}
//初始化线段树,至于为什么需要4倍空间,只要咱们理解了二叉树就一目了然了,线段树每个节点存储的就是某一段的区间和
func NewSegmentTree(num []int) *segmentTree {
countNum := len(num)
data := make([]int, countNum)
for k, v := range num {
data[k] = v
}
tree := make([]int, 4*countNum)
if countNum > 0 {
var buildTree func(int, int, int)
buildTree = func(index, left, right int) {
if left == right {
tree[index] = num[left]
return
}
leftChild := leftChild(index)
rightChild := rightChild(index)
mid := left + ((right - left) >> 1)
buildTree(leftChild, left, mid)
buildTree(rightChild, mid+1, right)
tree[index] = tree[leftChild] + tree[rightChild]
}
buildTree(0, 0, countNum-1)
}
return &segmentTree{data, tree}
}
然后就是求和操作了,我们把各个区间的和分别保存好了,本质上就变成了二叉树上找常数个节点,所以当然是O(logN)了:
//求和操作,只需要通过递归来找到最近的区间和就好
func (st *segmentTree) SumRange(start, end int) int {
var sum func(int, int, int, int, int) int
sum = func(index, left, right, start, end int) int {
if left == start && right == end {
return st.tree[index]
}
leftChild := leftChild(index)
rightChild := rightChild(index)
mid := left + ((right - left) >> 1)
if start >= mid+1 {
return sum(rightChild, mid+1, right, start, end)
} else if end <= mid {
return sum(leftChild, left, mid, start, end)
}
return sum(leftChild, left, mid, start, mid) + sum(rightChild, mid+1, right, mid+1, end)
}
return sum(0, 0, len(st.data)-1, start, end)
}
修改操作,好记性不如烂笔头,用笔画画就可以:
//修改操作,递归找到叶子节点改掉索引对应的值,然后回溯的过程改掉包含索引的所有区间的和
func (st *segmentTree) Update(i int, value int) {
countNum := len(st.data)
if i >= len(st.data) {
return
}
st.data[i] = value
var up func(int, int, int)
up = func(index, left, right int) {
if left == right {
st.tree[index] = value
return
}
leftChild := leftChild(index)
rightChild := rightChild(index)
mid := left + ((right - left) >> 1)
if i >= mid+1 {
up(rightChild, mid+1, right)
} else if i <= mid {
up(leftChild, left, mid)
}
st.tree[index] = st.tree[leftChild] + st.tree[rightChild]
}
up(0, 0, countNum-1)
}
func leftChild(i int) int {
return (i << 1) + 1
}
func rightChild(i int) int {
return (i << 1) + 2
}
新手的代码总是有很大的优化空间的,走过路过的大爷们要不吝赐教哦。
当然了,这仅仅只是一个简单的线段树,只能解决我的小需求。还有更多的方案,比如用链表来构建树肯定要更灵活,而且线段树并不仅仅用来求和,可以实现更多面向区间的操作,更好玩的是可以懒惰更新,每次修改不着急重新构建整棵树,可以一点一点地来,这样更加优化了性能。大家如果对线段树感兴趣的话,查资料去吧~o( ̄︶ ̄)o
算法梦想家,来跟我一起玩算法,玩音乐,聊聊文学创作,咱们一起天马行空!