golang源码分析：btree
golang源码分析：btree
github.com/google/btree是golang实现的一个平衡多叉树。它是etcd索引使用的树形结构。它使用起来非常简单。
package main
import (
"fmt"
"strconv"
"github.com/google/btree"
)
type MyTree struct {
Age int
Name string
}
func (m *MyTree) Less(item btree.Item) bool {
return m.Age < (item.(*MyTree)).Age
}
func main() {
tree := btree.New(2) //创建一个度为2的树，它最多有4个孩子节点和三个元素
for i := 0; i < 100; i++ {
//插入数据
tree.ReplaceOrInsert(&MyTree{Age: i, Name: "freedom" + strconv.Itoa(i)})
}
//查询数据
tree.DescendRange(&MyTree{Age: 52}, &MyTree{Age: 48}, func(a btree.Item) bool {
item := a.(*MyTree)
fmt.Println(item)
return true
})
}接着，我们复习下多叉平衡树的定义：对于一个度为m的多叉平衡树，最多有2*m+1个孩子节点和2*m个元素，所有非叶子非根节点至少有m个元素，所有叶子节点都在同一层级。
接着我们看下这个包里是如何实现的：它实现了两个版本go1.18以上的范型版本和go1.18 btree_generic.go以下的非范型版本btree.go。并且和左倾红黑树做了性能对比：btree_mem.go
这里我们重点看下非范型版本（范型版本逻辑一样，只是将一些基于接口实现的部分换成了范型实现），源码位于btree.go，如果要使用btree，需要实现Less接口即可。
type Item interface {
// Less tests whether the current item is less than the given argument.
//
// This must provide a strict weak ordering.
// If !a.Less(b) && !b.Less(a), we treat this to mean a == b (i.e. we can only
// hold one of either a or b in the tree).
Less(than Item) bool
}然后定义了FreeList里面保存了可以回收再利用的节点列表，用于减轻gc的压力
const (
DefaultFreeListSize = 32
)type FreeList struct {
mu sync.Mutex
freelist []*node
}所有新节点的分配和节点释放都是基于freeList来实现的，可以认为是一个简化的内存管理器：每次从节点slice的最后一个位置分配节点，缩小slice，如果slice为空就new 新创建一个节点。
func (f *FreeList) newNode() (n *node) {
f.mu.Lock()
index := len(f.freelist) - 1
if index < 0 {
f.mu.Unlock()
return new(node)
}
n = f.freelist[index]
f.freelist[index] = nil
f.freelist = f.freelist[:index]释放节点的过程刚好反过来，小于最大长度，把要释放的节点放在slice最后，超过长度不处理，让gc给处理掉。
func (f *FreeList) freeNode(n *node) (out bool) {
if len(f.freelist) < cap(f.freelist) {
f.freelist = append(f.freelist, n)
out = true它管理的node也就是我们的多叉树节点，其定义如下
type node struct {
items items //节点内元素
children children //孩子节点
cow *copyOnWriteContext //实现copy on write
}type items []Item
type children []*nodetype copyOnWriteContext struct {
freelist *FreeList
}func (c *copyOnWriteContext) newNode() (n *node) {
n = c.freelist.newNode()
n.cow = c
return
}func (c *copyOnWriteContext) freeNode(n *node) freeType {
if n.cow == c {
// clear to allow GC
n.items.truncate(0)
n.children.truncate(0)
n.cow = nil
if c.freelist.freeNode(n) {copyOnWriteContext也实现了节点分配和释放逻辑，分配的时候从freelist里面分配节点，并把当前context的地址赋值给新分配的节点。释放节点的时候，如果是相同context就释放，否则不操作。释放的时候先把元素和孩子节点指针清空，然后把节点放入当前context的freelist里面，放满了交给垃圾回收器处理。
通过逻辑保证孩子节点数永远比节点内元素个数多一个。BTree代表了一颗平衡多叉树，里面包含了这棵树的度和树里元素的总个数，是读并发安全的，但是不是写并发安全的。
type BTree struct {
degree int
length int
root *node
cow *copyOnWriteContext
}我们从New函数进入看下它是如何初始化的：
func New(degree int) *BTree {
return NewWithFreeList(degree, NewFreeList(DefaultFreeListSize))
}func NewWithFreeList(degree int, f *FreeList) *BTree {
return &BTree{
degree: degree,
cow: ©OnWriteContext{freelist: f},
}在cow里存储了预分配的节点指针空间
func NewFreeList(size int) *FreeList {
return &FreeList{freelist: make([]*node, 0, size)}
}在学习整棵树的增删改查之前，我们先看下节点内部元素的增删改查过程。
type items []Itemfunc (s *items) insertAt(index int, item Item) {
*s = append(*s, nil)
if index < len(*s) {
copy((*s)[index+1:], (*s)[index:])
}
(*s)[index] = item在slice的最后插入一个空指针，并把插入位置之后的所有元素后移一个位置，然后在空出的位置插入当前元素，删除是逆向过程，将当前位置之后的所有元素左移一位，最后位置置空
func (s *items) removeAt(index int) Item {
copy((*s)[index:], (*s)[index+1:])
(*s)[len(*s)-1] = nil返回最后的元素，并把最后的元素位置地址置空
func (s *items) pop() (out Item) {var (
nilItems = make(items, 16)
nilChildren = make(children, 16)
)把指定位置之后的所有指针16个一组全部置为空
func (s *items) truncate(index int) {
var toClear items
*s, toClear = (*s)[:index], (*s)[index:]
for len(toClear) > 0 {
toClear = toClear[copy(toClear, nilItems):]
}查找的过程就是slice的二分查找
func (s items) find(item Item) (index int, found bool) {
i := sort.Search(len(s), func(i int) bool {
return item.Less(s[i])
})
if i > 0 && !s[i-1].Less(item) {
return i - 1, true
}
return i, false对于孩子节点的操作是一模一样的
type children []*nodefunc (s *children) insertAt(index int, n *node) {
func (s *children) removeAt(index int) *node {
func (s *children) pop() (out *node) {
func (s *children) truncate(index int) {分析完节点的内部元素和孩子节点后，我们看下树的节点，为了便于节点的遍历，定义了迭代器，如果返回false，迭代器停止遍历，否则继续在树上遍历
type ItemIterator func(i Item) boolmutableFor用于判断当前上下文中，可以修改的节点，如果不能修改，就复制一份用来修改，从而实现了copy on write，复制的时候会复制内部元素和孩子节点指针
func (n *node) mutableFor(cow *copyOnWriteContext) *node {
if n.cow == cow {
return n
out := cow.newNode()
copy(out.items, n.items)
copy(out.children, n.children)获取某个孩子在当前上下文可以修改的节点。
func (n *node) mutableChild(i int) *node {
c := n.children[i].mutableFor(n.cow)
n.children[i] = c
return c
}把节点的元素和孩子节点在指定位置拆分成两部分，返回当前元素和拆分后的新节点
func (n *node) split(i int) (Item, *node) {
item := n.items[i]
next := n.cow.newNode()
next.items = append(next.items, n.items[i+1:]...)
n.items.truncate(i)
if len(n.children) > 0 {
next.children = append(next.children, n.children[i+1:]...)
n.children.truncate(i + 1)检查孩子节点是否需要拆分，如果孩子节点需要拆分，返回true
func (n *node) maybeSplitChild(i, maxItems int) bool {
if len(n.children[i].items) < maxItems {
return false
first := n.mutableChild(i)
item, second := first.split(maxItems / 2)
n.items.insertAt(i, item)
n.children.insertAt(i+1, second)找到i位置的孩子节点，并把孩子节点从最大节点数的一半地方拆开，在当前节点i的位置插入切分位置的元素暗，并在i+1的孩子节点位置插入从i位置孩子节点上拆分出来的右半部分。
接着看下节点的插入函数的执行流程，首先找到元素，如果找到了，替换这个位置的元素返回旧元素；没有找到执行插入新节点的流程。
插入的过程分两种情况：叶子节点（没有孩子节点的时候），直接在当前节点位置i插入新元素即可；非叶子节点如果需要拆分孩子节点拆分孩子节点，如果当前节点比插入点元素小，在i位置的孩子节点插入元素；否则在i+1位置的孩子节点插入元素。非叶子节点如果不需要拆分孩子节点，在孩子i的位置插入节点。插入过程是递归进行的。
func (n *node) insert(item Item, maxItems int) Item {
i, found := n.items.find(item)
if found {
out := n.items[i]
n.items[i] = item
return out
if len(n.children) == 0 {
n.items.insertAt(i, item)
if n.maybeSplitChild(i, maxItems) {
inTree := n.items[i]
switch {
case item.Less(inTree):
// no change, we want first split node
case inTree.Less(item):
i++ // we want second split node
default:
out := n.items[i]
n.items[i] = item
return out
}
}
return n.mutableChild(i).insert(item, maxItems)查找的过程是先通过二分查找方式在当前节点查找元素，如果找到了返回当前节点对应位置的元素，否则，在孩子i的位置查找元素，递归进行
func (n *node) get(key Item) Item {
i, found := n.items.find(key)
if found {
return n.items[i]
} else if len(n.children) > 0 {
return n.children[i].get(key)
}获取最小元素的时候，先从最小孩子的最小元素开始递归查找，如果没有孩子，返回当前节点的最小元素
func min(n *node) Item {
for len(n.children) > 0 {
n = n.children[0]
return n.items[0]最大元素查找是类似流程
func max(n *node) Item {
for len(n.children) > 0 {
n = n.children[len(n.children)-1]
return n.items[len(n.items)-1]删除的过程比较复杂：1，如果是叶子节点，并且找到了这个元素，直接删除即可；2，如果i位置的孩子节点元素个数比最小个数少，调用函数growChildAndRemove；3，如果在孩子节点i处找到了元素，移除孩子节点i的最大元素，并将它填入当前节点的i位置；4，如果没有找到，递归在孩子节点i处继续查找。
func (n *node) remove(item Item, minItems int, typ toRemove) Item {
switch typ {
case removeMax:
if len(n.children) == 0 {
return n.items.pop()
}
i = len(n.items)
case removeMin:
if len(n.children) == 0 {
return n.items.removeAt(0)
}
i = 0
case removeItem:
i, found = n.items.find(item)
if len(n.children) == 0 {
if found {
return n.items.removeAt(i)
}
return nil
}
default:
panic("invalid type")
}
if len(n.children[i].items) <= minItems {
return n.growChildAndRemove(i, item, minItems, typ)
}
child := n.mutableChild(i)
if found {
// The item exists at index 'i', and the child we've selected can give us a
// predecessor, since if we've gotten here it's got > minItems items in it.
out := n.items[i]
// We use our special-case 'remove' call with typ=maxItem to pull the
// predecessor of item i (the rightmost leaf of our immediate left child)
// and set it into where we pulled the item from.
n.items[i] = child.remove(nil, minItems, removeMax)
return out
}
return child.remove(item, minItems, typ)分三种场景进行调整，调整完毕后再递归调用remove方法来进行节点删除：1，孩子的左兄弟节点位置的元素个数比最小元素个数大，在孩子的左兄弟节点偷取最大元素，存入i-1位置，并把i-1位置的元素存入孩子i的最左侧，并把被偷位置的孩子节点插入到孩子i的最左孩子位置。2，如果孩子的右兄弟元素个数比最小元素个数大，从孩子的右兄弟的最左孩子位置处偷元素。插入到孩子i的最右位置。
func (n *node) growChildAndRemove(i int, item Item, minItems int, typ toRemove) Item {
if i > 0 && len(n.children[i-1].items) > minItems {
// Steal from left child
child := n.mutableChild(i)
stealFrom := n.mutableChild(i - 1)
stolenItem := stealFrom.items.pop()
child.items.insertAt(0, n.items[i-1])
n.items[i-1] = stolenItem
if len(stealFrom.children) > 0 {
child.children.insertAt(0, stealFrom.children.pop())
}
} else if i < len(n.items) && len(n.children[i+1].items) > minItems {
// steal from right child
child := n.mutableChild(i)
stealFrom := n.mutableChild(i + 1)
stolenItem := stealFrom.items.removeAt(0)
child.items = append(child.items, n.items[i])
n.items[i] = stolenItem
if len(stealFrom.children) > 0 {
child.children = append(child.children, stealFrom.children.removeAt(0))
}
} else {
if i >= len(n.items) {
i--
}
child := n.mutableChild(i)
// merge with right child
mergeItem := n.items.removeAt(i)
mergeChild := n.children.removeAt(i + 1)
child.items = append(child.items, mergeItem)
child.items = append(child.items, mergeChild.items...)
child.children = append(child.children, mergeChild.children...)
n.cow.freeNode(mergeChild)
}
return n.remove(item, minItems, typ)3，如果不满足上述两种情况，先把i位置的元素删掉，然后把i+1孩子给删掉，把删掉的元素和删掉的孩子的元素追加到孩子i的元素末尾，把删掉的孩子的孩子，追加到孩子i的孩子末尾。处理完上述调整后，就继续进行递归的删除操作。
iterate实现了迭代器，在节点内通过二分查找快速定位元素，然后再树形父子结构上通过递归来定位元素，迭代的时候需要控制方向的正反。
func (n *node) iterate(dir direction, start, stop Item, includeStart bool, hit bool, iter ItemIterator) (bool, bool) {
switch dir {
case ascend:
if start != nil {
index, _ = n.items.find(start)
}
for i := index; i < len(n.items); i++ {
if len(n.children) > 0 {
if hit, ok = n.children[i].iterate(dir, start, stop, includeStart, hit, iter); !ok {
return hit, false
}学习完node的增删改查算法后，BTree的增删改查只是对上述方法的一个包装而已。
Clone方法，通过对cow属性的简单复制，实现了懒复制。
func (t *BTree) Clone() (t2 *BTree) {
cow1, cow2 := *t.cow, *t.cow
out := *t
t.cow = &cow1
out.cow = &cow2
return &out 定义了节点的最大元素个数和最小元素个数。
func (t *BTree) maxItems() int {
return t.degree*2 - 1func (t *BTree) minItems() int {
return t.degree - 1插入元素过程，如果根节点为空，直接插入元素；如果根节点的元素个数达到最大了，将根节点一份为二，新生成一个节点作为它们的根节点；在根节点处插入元素。
func (t *BTree) ReplaceOrInsert(item Item) Item {
if t.root == nil {
t.root = t.cow.newNode()
t.root.items = append(t.root.items, item)
t.length++
return nil
t.root = t.root.mutableFor(t.cow)
if len(t.root.items) >= t.maxItems() {
item2, second := t.root.split(t.maxItems() / 2)
out := t.root.insert(item, t.maxItems())func (t *BTree) Delete(item Item) Item {
return t.deleteItem(item, removeItem)删除节点的过程是，首先移除元素，如果移除后根节点为空，把第一个孩子节点当作根节点，删除老的根节点。
func (t *BTree) deleteItem(item Item, typ toRemove) Item {
t.root = t.root.mutableFor(t.cow)
out := t.root.remove(item, t.minItems(), typ)
if len(t.root.items) == 0 && len(t.root.children) > 0 {
oldroot := t.root
t.root = t.root.children[0]
t.cow.freeNode(oldroot)
}range操作，就是迭代器的执行
func (t *BTree) AscendRange(greaterOrEqual, lessThan Item, iterator ItemIterator) {
t.root.iterate(ascend, greaterOrEqual, lessThan, true, false, iterator)func (t *BTree) Ascend(iterator ItemIterator) {func (t *BTree) DescendRange(lessOrEqual, greaterThan Item, iterator ItemIterator) {查找类似
func (t *BTree) Get(key Item) Item {
return t.root.get(key)func (t *BTree) Clear(addNodesToFreelist bool) {
t.root.reset(t.cow) func (n *node) reset(c *copyOnWriteContext) bool {
for _, child := range n.children {
if !child.reset(c) {BTree实现了b树，并通过cow属性， cow相等才能修改，否则只能复制 来实现了copy on write；通过编译选项，来控制使用范型版本还是非范型版本：
//go:build !go1.18
// +build !go1.18至此整个BTree的实现分析完毕。
本文分享自 golang算法架构leetcode技术php 微信公众号，前往查看
如有侵权，请联系 cloudcommunity@tencent.com 删除。
本文参与 腾讯云自媒体分享计划 ，欢迎热爱写作的你一起参与！