B+树Go实现

B+树动画演示
https://www.cs.usfca.edu/~galles/visualization/BPlusTree.html

package main

/*
B+树动画演示
https://www.cs.usfca.edu/~galles/visualization/BPlusTree.html
*/
import (
	"fmt"
)

const (
	MAX_KEYS = 3 // 每个节点的最大键数量
)

type BPlusTree struct {
	root *Node
}

type Node struct {
	keys     []int
	leaf     bool
	children []*Node
}

// 创建一个新的节点
func newNode(leaf bool) *Node {
	return &Node{
		keys:     make([]int, 0),
		leaf:     leaf,
		children: make([]*Node, 0),
	}
}

// 初始化B+树
func NewBPlusTree() *BPlusTree {
	root := newNode(true)
	return &BPlusTree{root: root}
}
func (tree *BPlusTree) PrintTree() {
	fmt.Println("B+ 树结构:")
	printNode(tree.root, 0, "", -1)
}

func printNode(node *Node, level int, prefix string, childNum int) {
	if childNum >= 0 {
		fmt.Printf("%s子节点 %d: ", prefix, childNum)
	} else {
		fmt.Printf("%s", prefix)
	}
	fmt.Printf("层级 %d: ", level)
	for i, k := range node.keys {
		fmt.Printf("%d", k)
		if i < len(node.keys)-1 {
			fmt.Print(" | ")
		}
	}
	fmt.Println()

	if !node.leaf {
		childPrefix := prefix + "│   "
		for i, child := range node.children {
			if i == len(node.children)-1 {
				fmt.Printf("%s└───", prefix)
				childPrefix = prefix + "    "
			} else {
				fmt.Printf("%s├───", prefix)
			}
			printNode(child, level+1, childPrefix, i)
		}
	}
}

func (tree *BPlusTree) Search(key int) bool {
	fmt.Printf("搜索键值 %d:\n", key)
	return searchNode(tree.root, key, 0, "", -1)
}

func searchNode(node *Node, key int, level int, prefix string, childNum int) bool {
	if childNum >= 0 {
		fmt.Printf("%s子节点 %d: ", prefix, childNum)
	} else {
		fmt.Printf("%s", prefix)
	}
	fmt.Printf("层级 %d: 检查节点，键值为 [", level)
	for i, k := range node.keys {
		fmt.Printf("%d", k)
		if i < len(node.keys)-1 {
			fmt.Print(" | ")
		}
	}
	fmt.Println("]")

	if node.leaf {
		for _, k := range node.keys {
			if k == key {
				fmt.Printf("%s  在叶子节点中找到键值 %d\n", prefix, key)
				return true
			}
		}
		fmt.Printf("%s  在叶子节点中未找到键值 %d\n", prefix, key)
		return false
	}

	for i := 0; i < len(node.keys); i++ {
		if key < node.keys[i] {
			fmt.Printf("%s  键值 %d < %d，移动到子节点 %d\n", prefix, key, node.keys[i], i)
			return searchNode(node.children[i], key, level+1, prefix+"    ", i)
		}
	}
	lastChild := len(node.keys)
	fmt.Printf("%s  键值 %d >= 节点中所有键值，移动到子节点 %d\n", prefix, key, lastChild)
	return searchNode(node.children[lastChild], key, level+1, prefix+"    ", lastChild)
}

// 插入一个键
func (tree *BPlusTree) Insert(key int) {
	root := tree.root
	if len(root.keys) == MAX_KEYS {
		newRoot := newNode(false)
		newRoot.children = append(newRoot.children, root)
		tree.splitChild(newRoot, 0)
		tree.root = newRoot
	}

	tree.insertNonFull(root, key)
}

func (tree *BPlusTree) insertNonFull(node *Node, key int) {
	if node.leaf {
		node.keys = append(node.keys, key)
		sortKeys(node)
		return
	}

	// 找到插入位置
	i := len(node.keys) - 1
	for i >= 0 && key < node.keys[i] {
		i--
	}
	i++

	// 插入到子节点
	child := node.children[i]
	if len(child.keys) == MAX_KEYS {
		tree.splitChild(node, i)
		if key > node.keys[i] {
			i++
		}
	}
	tree.insertNonFull(node.children[i], key)
}

func (tree *BPlusTree) splitChild(parent *Node, index int) {
	node := parent.children[index]
	newNode := newNode(node.leaf)

	// 将中间的键提升到父节点
	mid := len(node.keys) / 2
	parent.keys = append(parent.keys, 0)
	copy(parent.keys[index+1:], parent.keys[index:])
	parent.keys[index] = node.keys[mid]

	// 将子节点分割
	parent.children = append(parent.children, nil)
	copy(parent.children[index+2:], parent.children[index+1:])
	parent.children[index+1] = newNode

	// 将右半部分的键和子节点转移到新的节点
	newNode.keys = append(newNode.keys, node.keys[mid+1:]...)
	node.keys = node.keys[:mid]
	if !node.leaf {
		newNode.children = append(newNode.children, node.children[mid+1:]...)
		node.children = node.children[:mid+1]
	}
}

// 排序节点的键
func sortKeys(node *Node) {
	for i := 1; i < len(node.keys); i++ {
		for j := i; j > 0 && node.keys[j] < node.keys[j-1]; j-- {
			node.keys[j], node.keys[j-1] = node.keys[j-1], node.keys[j]
		}
	}
}

func main() {
	tree := NewBPlusTree()

	keys := []int{10, 20, 5, 6, 15, 30, 40, 50, 60}
	for _, key := range keys {
		tree.Insert(key)
	}

	fmt.Println("B+ Tree structure:")
	tree.PrintTree()

	fmt.Println("\nSearch operations:")
	searchKeys := []int{15, 25, 40, 7}
	for _, key := range searchKeys {
		if tree.Search(key) {
			fmt.Printf("Found %d\n", key)
		} else {
			fmt.Printf("%d not found\n", key)
		}
		fmt.Println()
	}
}