diff --git a/alg/bplus.go b/alg/bplus.go new file mode 100644 index 0000000..79e67db --- /dev/null +++ b/alg/bplus.go @@ -0,0 +1,355 @@ +package alg + +import "strings" + +type IndexKey interface { + Compare(other IndexKey) int +} + +type SortOrder struct { + Field string + Desc bool +} + +type IndexEntry struct { + Key IndexKey + Value interface{} +} + +const BTreeMinDegree = 2 + +type BTreeNode struct { + Keys []IndexKey + Values []interface{} // 叶子节点存储数据 + Childs []*BTreeNode // 非叶子节点存储子节点 + IsLeaf bool + Next *BTreeNode // 叶子节点之间的链表指针 + Prev *BTreeNode +} + +type BTree struct { + Root *BTreeNode +} + +func NewBTreeNode(leaf bool) *BTreeNode { + return &BTreeNode{ + Keys: make([]IndexKey, 0), + Values: make([]interface{}, 0), + Childs: make([]*BTreeNode, 0), + IsLeaf: leaf, + } +} + +func NewBTree() *BTree { + return &BTree{ + Root: NewBTreeNode(true), + } +} + +func (bt *BTree) Insert(key IndexKey, value interface{}) { + root := bt.Root + // 如果根节点已满,需要分裂 + if len(root.Keys) == 2*BTreeMinDegree-1 { + newRoot := NewBTreeNode(false) + newRoot.Childs = append(newRoot.Childs, root) + splitChild(newRoot, 0) + bt.Root = newRoot + } + insertNonFull(bt.Root, key, value) +} + +// 分裂子节点 +func splitChild(parent *BTreeNode, index int) { + minDegree := BTreeMinDegree + fullNode := parent.Childs[index] + // 创建新节点 + newNode := NewBTreeNode(fullNode.IsLeaf) + // 中间键位置 + mid := minDegree - 1 + + // 拆分键和值 + newNode.Keys = fullNode.Keys[mid:] + fullNode.Keys = fullNode.Keys[:mid] + + if !fullNode.IsLeaf { + newNode.Childs = fullNode.Childs[mid:] + fullNode.Childs = fullNode.Childs[:mid] + } else { + // 叶子节点需要维护链表 + newNode.Values = fullNode.Values[mid:] + fullNode.Values = fullNode.Values[:mid] + newNode.Next = fullNode.Next + newNode.Prev = fullNode + if fullNode.Next != nil { + fullNode.Next.Prev = newNode + } + fullNode.Next = newNode + } + + // 将中间键提升到父节点 + parent.Keys = append(parent.Keys[:index+1], append([]IndexKey{fullNode.Keys[mid]}, parent.Keys[index+1:]...)...) + fullNode.Keys = fullNode.Keys[:mid] +} + +// 插入非满节点 +func insertNonFull(node *BTreeNode, key IndexKey, value interface{}) { + i := len(node.Keys) - 1 + + if node.IsLeaf { + // 在叶子节点插入 + for i >= 0 && key.Compare(node.Keys[i]) < 0 { + i-- + } + // 插入键值 + if i >= 0 { + node.Keys = append(node.Keys[:i+1], append([]IndexKey{key}, node.Keys[i+1:]...)...) + node.Values = append(node.Values[:i+1], append([]interface{}{value}, node.Values[i+1:]...)...) + } else { + node.Keys = append([]IndexKey{key}, node.Keys...) + node.Values = append([]interface{}{value}, node.Values...) + } + } else { + // 找到子节点 + for i >= 0 && key.Compare(node.Keys[i]) < 0 { + i-- + } + i++ + + // 如果子节点已满,需要分裂 + child := node.Childs[i] + if len(child.Keys) == 2*BTreeMinDegree-1 { + splitChild(node, i) + if key.Compare(node.Keys[i]) > 0 { + i++ // 选择正确的位置 + } + } + insertNonFull(node.Childs[i], key, value) + } +} + +func (bt *BTree) Search(key IndexKey) (interface{}, bool) { + node := bt.Root + for node != nil { + // 在当前节点查找位置 + i := 0 + for i < len(node.Keys) && key.Compare(node.Keys[i]) > 0 { + i++ + } + + if node.IsLeaf { + // 在叶子节点中查找精确匹配 + if i < len(node.Keys) && key.Compare(node.Keys[i]) == 0 { + return node.Values[i], true + } + return nil, false + } else { + // 继续查找子节点 + node = node.Childs[i] + } + } + return nil, false +} + +func (bt *BTree) Delete(key IndexKey) { + deleteFromNode(bt.Root, key) +} + +// 从节点删除指定键 +func deleteFromNode(node *BTreeNode, key IndexKey) { + minDegree := BTreeMinDegree + // 寻找键的位置 + i := 0 + for i < len(node.Keys) && key.Compare(node.Keys[i]) > 0 { + i++ + } + + if node.IsLeaf { + // 叶子节点直接删除 + if i < len(node.Keys) && key.Compare(node.Keys[i]) == 0 { + // 删除键值对 + node.Keys = append(node.Keys[:i], node.Keys[i+1:]...) + node.Values = append(node.Values[:i], node.Values[i+1:]...) + } + } else { + // 非叶子节点 + if i < len(node.Keys) && key.Compare(node.Keys[i]) == 0 { + // 找到键,需要替换为后继值 + child := node.Childs[i] + sibling := node.Childs[i+1] + + if len(child.Keys) >= minDegree { + // 前驱节点存在,替换为前驱 + predecessor := getPredecessor(child) + node.Keys[i] = predecessor + deleteFromNode(child, predecessor) + } else if len(sibling.Keys) >= minDegree { + // 后继节点存在,替换为后继 + successor := getSuccessor(sibling) + node.Keys[i] = successor + deleteFromNode(sibling, successor) + } else { + // 合并节点 + mergeNodes(node, i) + deleteFromNode(child, key) + } + } else { + // 继续删除子节点 + child := node.Childs[i] + if len(child.Keys) >= minDegree { + // 子节点足够大,继续删除 + deleteFromNode(child, key) + } else { + // 需要重新平衡 + if i > 0 && len(node.Childs[i-1].Keys) >= minDegree { + // 从左兄弟借 + borrowFromLeft(node, i) + } else if i < len(node.Childs)-1 && len(node.Childs[i+1].Keys) >= minDegree { + // 从右兄弟借 + borrowFromRight(node, i) + } else { + // 合并节点 + if i > 0 { + mergeNodes(node, i-1) + } else { + mergeNodes(node, i) + } + } + deleteFromNode(child, key) + } + } + } +} + +// 获取前驱节点 +func getPredecessor(node *BTreeNode) IndexKey { + for !node.IsLeaf { + node = node.Childs[len(node.Childs)-1] + } + return node.Keys[len(node.Keys)-1] +} + +// 获取后继节点 +func getSuccessor(node *BTreeNode) IndexKey { + for !node.IsLeaf { + node = node.Childs[0] + } + return node.Keys[0] +} + +// 合并节点 +func mergeNodes(parent *BTreeNode, index int) { + left := parent.Childs[index] + right := parent.Childs[index+1] + + // 将左节点和右节点合并 + left.Keys = append(left.Keys, parent.Keys[index]) + left.Keys = append(left.Keys, right.Keys...) + + if left.IsLeaf { + left.Values = append(left.Values, right.Values...) + left.Next = right.Next + if right.Next != nil { + right.Next.Prev = left + } + } else { + left.Childs = append(left.Childs, right.Childs...) + } + + // 更新父节点 + parent.Keys = append(parent.Keys[:index], parent.Keys[index+1:]...) + parent.Childs = append(parent.Childs[:index+1], parent.Childs[index+2:]...) +} + +// 从左兄弟借 +func borrowFromLeft(parent *BTreeNode, index int) { + left := parent.Childs[index-1] + current := parent.Childs[index] + + // 移动键 + current.Keys = append([]IndexKey{parent.Keys[index-1]}, current.Keys...) + if !left.IsLeaf { + current.Childs = append([]*BTreeNode{left.Childs[len(left.Childs)-1]}, current.Childs...) + left.Childs = left.Childs[:len(left.Childs)-1] + } + + // 更新父节点 + parent.Keys[index-1] = left.Keys[len(left.Keys)-1] + left.Keys = left.Keys[:len(left.Keys)-1] +} + +// 从右兄弟借 +func borrowFromRight(parent *BTreeNode, index int) { + current := parent.Childs[index] + right := parent.Childs[index+1] + + // 移动键 + current.Keys = append(current.Keys, parent.Keys[index]) + parent.Keys[index] = right.Keys[0] + + if !current.IsLeaf { + current.Childs = append(current.Childs, right.Childs[0]) + right.Childs = right.Childs[1:] + } + + right.Keys = right.Keys[1:] +} + +func (bt *BTree) Traverse(cb func(key IndexKey) bool) []interface{} { + result := make([]interface{}, 0) + // 找到最左边的叶子节点 + node := bt.Root + for node != nil && !node.IsLeaf { + node = node.Childs[0] + } + + // 遍历所有叶子节点 + for node != nil { + for i, key := range node.Keys { + if cb(key) { + result = append(result, node.Values[i]) + } else { + return result + } + } + node = node.Next + } + return result +} + +type MultiFieldKey struct { + Fields []string + SortDirs []bool // true 表示降序 + Values map[string]interface{} +} + +func (m MultiFieldKey) Compare(other IndexKey) int { + o := other.(MultiFieldKey) + for i, field := range m.Fields { + v1 := m.Values[field] + v2 := o.Values[field] + + var cmp int + switch v1.(type) { + case string: + cmp = strings.Compare(v1.(string), v2.(string)) + case float64: + if v1.(float64) < v2.(float64) { + cmp = -1 + } else if v1.(float64) > v2.(float64) { + cmp = 1 + } else { + cmp = 0 + } + default: + panic("Unsupported type") + } + + if cmp != 0 { + if m.SortDirs[i] { + return -cmp + } + return cmp + } + } + return 0 +}