add bplus alg
This commit is contained in:
parent
00ec4fd001
commit
fd28639f17
|
@ -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
|
||||
}
|
Loading…
Reference in New Issue