手寫線段樹 第一場第一鏡

靜謐的夜最適合刷算法題了。刷着刷着發現了一個好玩的數據結構叫作線段樹，聽說是算法競賽的常客哦，因而就本身寫一個簡單的玩玩。

原理

想了解線段樹的同志們請移步，裏邊有原理和示意圖： 百度百科，線段樹

需求和痛點

我用它來主要是爲了快速找到數組某區間內的數字和，而且在修改數組某幾個元素以後再次找區間內的數字和。可想而知，我有兩個需求：求和，修改。

正常狀況下，由於咱們並不想給數組排序，那麼，咱們能夠用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

算法夢想家，來跟我一塊兒玩算法，玩音樂，聊聊文學創做，我們一塊兒天馬行空！