手写线段树 第一场第一镜

静谧的夜最适合刷算法题了。刷着刷着发现了一个好玩的数据结构叫做线段树，据说是算法竞赛的常客哦，于是就自己写一个简单的玩玩。

原理

想了解线段树的同志们请移步，里边有原理和示意图： 百度百科，线段树

需求和痛点

我用它来主要是为了快速找到数组某区间内的数字和，并且在修改数组某几个元素之后再次找区间内的数字和。可想而知，我有两个需求：求和，修改。

正常情况下，因为我们并不想给数组排序，那么，我们可以用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

算法梦想家，来跟我一起玩算法，玩音乐，聊聊文学创作，咱们一起天马行空！

转载于:https://juejin.im/post/5d0b0cd3e51d4510a73280cb