diff --git a/g/container/gtree/gtree_btree.go b/g/container/gtree/gtree_btree.go new file mode 100644 index 000000000..fa0d30fa8 --- /dev/null +++ b/g/container/gtree/gtree_btree.go @@ -0,0 +1,830 @@ +// Copyright 2019 gf Author(https://github.com/gogf/gf). All Rights Reserved. +// +// This Source Code Form is subject to the terms of the MIT License. +// If a copy of the MIT was not distributed with this file, +// You can obtain one at https://github.com/gogf/gf. + +package gtree + +import ( + "bytes" + "fmt" + "github.com/gogf/gf/g/container/gvar" + "github.com/gogf/gf/g/internal/rwmutex" + "strings" +) + +// BTree holds elements of the B-tree. +type BTree struct { + mu *rwmutex.RWMutex + root *BTreeNode + comparator func(v1, v2 interface{}) int + size int // Total number of keys in the tree + m int // order (maximum number of children) +} + +// BTreeNode is a single element within the tree. +type BTreeNode struct { + Parent *BTreeNode + Entries []*BTreeNodeEntry // Contained keys in node + Children []*BTreeNode // Children nodes +} + +// BTreeNodeEntry represents the key-value pair contained within nodes. +type BTreeNodeEntry struct { + Key interface{} + Value interface{} +} + +// NewBTree instantiates a B-tree with (maximum number of children) and a custom key comparator. +// Note that the must be greater or equal than 3, or else it panics. +func NewBTree(m int, comparator func(v1, v2 interface{}) int, unsafe...bool) *BTree { + if m < 3 { + panic("Invalid order, should be at least 3") + } + return &BTree{ + comparator : comparator, + mu : rwmutex.New(unsafe...), + m : m, + } +} + +// Clone returns a new tree with a copy of current tree. +func (tree *BTree) Clone(unsafe ...bool) *BTree { + newTree := NewBTree(tree.m, tree.comparator, !tree.mu.IsSafe()) + newTree.Sets(tree.Map()) + return newTree +} + +// Set inserts key-value item into the tree. +func (tree *BTree) Set(key interface{}, value interface{}) { + tree.mu.Lock() + defer tree.mu.Unlock() + tree.doSet(key, value) +} + +// doSet inserts key-value pair node into the tree. +// If key already exists, then its value is updated with the new value. +func (tree *BTree) doSet(key interface{}, value interface{}) { + entry := &BTreeNodeEntry{Key: key, Value: value} + if tree.root == nil { + tree.root = &BTreeNode{Entries: []*BTreeNodeEntry{entry}, Children: []*BTreeNode{}} + tree.size++ + return + } + + if tree.insert(tree.root, entry) { + tree.size++ + } +} + +// Sets batch sets key-values to the tree. +func (tree *BTree) Sets(data map[interface{}]interface{}) { + tree.mu.Lock() + defer tree.mu.Unlock() + for k, v := range data { + tree.doSet(k, v) + } +} + +// Get searches the node in the tree by and returns its value or nil if key is not found in tree. +func (tree *BTree) Get(key interface{}) (value interface{}) { + value, _ = tree.Search(key) + return +} + +// doSetWithLockCheck checks whether value of the key exists with mutex.Lock, +// if not exists, set value to the map with given , +// or else just return the existing value. +// +// When setting value, if is type of , +// it will be executed with mutex.Lock of the hash map, +// and its return value will be set to the map with . +// +// It returns value with given . +func (tree *BTree) doSetWithLockCheck(key interface{}, value interface{}) interface{} { + tree.mu.Lock() + defer tree.mu.Unlock() + if entry := tree.doSearch(key); entry != nil { + return entry.Value + } + if f, ok := value.(func() interface {}); ok { + value = f() + } + tree.doSet(key, value) + return value +} + +// GetOrSet returns the value by key, +// or set value with given if not exist and returns this value. +func (tree *BTree) GetOrSet(key interface{}, value interface{}) interface{} { + if v, ok := tree.Search(key); !ok { + return tree.doSetWithLockCheck(key, value) + } else { + return v + } +} + +// GetOrSetFunc returns the value by key, +// or sets value with return value of callback function if not exist +// and returns this value. +func (tree *BTree) GetOrSetFunc(key interface{}, f func() interface{}) interface{} { + if v, ok := tree.Search(key); !ok { + return tree.doSetWithLockCheck(key, f()) + } else { + return v + } +} + +// GetOrSetFuncLock returns the value by key, +// or sets value with return value of callback function if not exist +// and returns this value. +// +// GetOrSetFuncLock differs with GetOrSetFunc function is that it executes function +// with mutex.Lock of the hash map. +func (tree *BTree) GetOrSetFuncLock(key interface{}, f func() interface{}) interface{} { + if v, ok := tree.Search(key); !ok { + return tree.doSetWithLockCheck(key, f) + } else { + return v + } +} + +// GetVar returns a gvar.Var with the value by given . +// The returned gvar.Var is un-concurrent safe. +func (tree *BTree) GetVar(key interface{}) *gvar.Var { + return gvar.New(tree.Get(key), true) +} + +// GetVarOrSet returns a gvar.Var with result from GetVarOrSet. +// The returned gvar.Var is un-concurrent safe. +func (tree *BTree) GetVarOrSet(key interface{}, value interface{}) *gvar.Var { + return gvar.New(tree.GetOrSet(key, value), true) +} + +// GetVarOrSetFunc returns a gvar.Var with result from GetOrSetFunc. +// The returned gvar.Var is un-concurrent safe. +func (tree *BTree) GetVarOrSetFunc(key interface{}, f func() interface{}) *gvar.Var { + return gvar.New(tree.GetOrSetFunc(key, f), true) +} + +// GetVarOrSetFuncLock returns a gvar.Var with result from GetOrSetFuncLock. +// The returned gvar.Var is un-concurrent safe. +func (tree *BTree) GetVarOrSetFuncLock(key interface{}, f func() interface{}) *gvar.Var { + return gvar.New(tree.GetOrSetFuncLock(key, f), true) +} + +// SetIfNotExist sets to the map if the does not exist, then return true. +// It returns false if exists, and would be ignored. +func (tree *BTree) SetIfNotExist(key interface{}, value interface{}) bool { + if !tree.Contains(key) { + tree.doSetWithLockCheck(key, value) + return true + } + return false +} + +// SetIfNotExistFunc sets value with return value of callback function , then return true. +// It returns false if exists, and would be ignored. +func (tree *BTree) SetIfNotExistFunc(key interface{}, f func() interface{}) bool { + if !tree.Contains(key) { + tree.doSetWithLockCheck(key, f()) + return true + } + return false +} + +// SetIfNotExistFuncLock sets value with return value of callback function , then return true. +// It returns false if exists, and would be ignored. +// +// SetIfNotExistFuncLock differs with SetIfNotExistFunc function is that +// it executes function with mutex.Lock of the hash map. +func (tree *BTree) SetIfNotExistFuncLock(key interface{}, f func() interface{}) bool { + if !tree.Contains(key) { + tree.doSetWithLockCheck(key, f) + return true + } + return false +} + +// Contains checks whether exists in the tree. +func (tree *BTree) Contains(key interface{}) bool { + _, ok := tree.Search(key) + return ok +} + +// Remove remove the node from the tree by key. +// Key should adhere to the comparator's type assertion, otherwise method panics. +func (tree *BTree) doRemove(key interface{}) (value interface{}) { + node, index, found := tree.searchRecursively(tree.root, key) + if found { + value = node.Entries[index].Value + tree.delete(node, index) + tree.size-- + } + return +} + +// Remove removes the node from the tree by . +func (tree *BTree) Remove(key interface{}) (value interface{}) { + tree.mu.Lock() + defer tree.mu.Unlock() + return tree.doRemove(key) +} + +// Removes batch deletes values of the tree by . +func (tree *BTree) Removes(keys []interface{}) { + tree.mu.Lock() + defer tree.mu.Unlock() + for key := range keys { + tree.doRemove(key) + } +} + +// Empty returns true if tree does not contain any nodes +func (tree *BTree) IsEmpty() bool { + return tree.Size() == 0 +} + +// Size returns number of nodes in the tree. +func (tree *BTree) Size() int { + tree.mu.RLock() + defer tree.mu.RUnlock() + return tree.size +} + +// Keys returns all keys in asc order. +func (tree *BTree) Keys() []interface{} { + keys := make([]interface{}, tree.Size()) + index := 0 + tree.IteratorAsc(func(key, value interface{}) bool { + keys[index] = key + index++ + return true + }) + return keys +} + +// Values returns all values in asc order based on the key. +func (tree *BTree) Values() []interface{} { + values := make([]interface{}, tree.Size()) + index := 0 + tree.IteratorAsc(func(key, value interface{}) bool { + values[index] = key + index++ + return true + }) + return values +} + +// Map returns all key-value items as map. +func (tree *BTree) Map() map[interface{}]interface{} { + m := make(map[interface{}]interface{}, tree.Size()) + tree.IteratorAsc(func(key, value interface{}) bool { + m[key] = value + return true + }) + return m +} + +// Clear removes all nodes from the tree. +func (tree *BTree) Clear() { + tree.mu.Lock() + defer tree.mu.Unlock() + tree.root = nil + tree.size = 0 +} + +// Height returns the height of the tree. +func (tree *BTree) Height() int { + tree.mu.RLock() + defer tree.mu.RUnlock() + return tree.root.height() +} + +// Left returns the left-most (min) entry or nil if tree is empty. +func (tree *BTree) Left() *BTreeNodeEntry { + tree.mu.RLock() + defer tree.mu.RUnlock() + node := tree.left(tree.root) + return node.Entries[0] +} + +// Right returns the right-most (max) entry or nil if tree is empty. +func (tree *BTree) Right() *BTreeNodeEntry { + tree.mu.RLock() + defer tree.mu.RUnlock() + node := tree.right(tree.root) + return node.Entries[len(node.Entries) - 1] +} + +// String returns a string representation of container (for debugging purposes) +func (tree *BTree) String() string { + tree.mu.RLock() + defer tree.mu.RUnlock() + var buffer bytes.Buffer + if _, err := buffer.WriteString("BTree\n"); err != nil { + } + if tree.size != 0 { + tree.output(&buffer, tree.root, 0, true) + } + return buffer.String() +} + +// Search searches the tree with given . +// Second return parameter is true if key was found, otherwise false. +func (tree *BTree) Search(key interface{}) (value interface{}, found bool) { + tree.mu.RLock() + defer tree.mu.RUnlock() + node, index, found := tree.searchRecursively(tree.root, key) + if found { + return node.Entries[index].Value, true + } + return nil, false +} + +// Search searches the tree with given without mutex. +// It returns the entry if found or otherwise nil. +func (tree *BTree) doSearch(key interface{}) *BTreeNodeEntry { + node, index, found := tree.searchRecursively(tree.root, key) + if found { + return node.Entries[index] + } + return nil +} + +// Print prints the tree to stdout. +func (tree *BTree) Print() { + fmt.Println(tree.String()) +} + +func (entry *BTreeNodeEntry) String() string { + return fmt.Sprintf("%v", entry.Key) +} + +// IteratorAsc iterates the tree in ascending order with given callback function . +// If returns true, then it continues iterating; or false to stop. +func (tree *BTree) IteratorAsc(f func (key, value interface{}) bool) { + tree.mu.RLock() + defer tree.mu.RUnlock() + node := tree.left(tree.root) + if node == nil { + return + } + entry := node.Entries[0] +loop: + if entry == nil { + return + } + if !f(entry.Key, entry.Value) { + return + } + // Find current entry position in current node + e, _ := tree.search(node, entry.Key) + // Try to go down to the child right of the current entry + if e + 1 < len(node.Children) { + node = node.Children[e + 1] + // Try to go down to the child left of the current node + for len(node.Children) > 0 { + node = node.Children[0] + } + // Return the left-most entry + entry = node.Entries[0] + goto loop + } + // Above assures that we have reached a leaf node, so return the next entry in current node (if any) + if e + 1 < len(node.Entries) { + entry = node.Entries[e + 1] + goto loop + } + // Reached leaf node and there are no entries to the right of the current entry, so go up to the parent + for node.Parent != nil { + node = node.Parent + // Find next entry position in current node (note: search returns the first equal or bigger than entry) + e, _ := tree.search(node, entry.Key) + // Check that there is a next entry position in current node + if e < len(node.Entries) { + entry = node.Entries[e] + goto loop + } + } +} + +// IteratorDesc iterates the tree in descending order with given callback function . +// If returns true, then it continues iterating; or false to stop. +func (tree *BTree) IteratorDesc(f func (key, value interface{}) bool) { + tree.mu.RLock() + defer tree.mu.RUnlock() + node := tree.right(tree.root) + if node == nil { + return + } + entry := node.Entries[len(node.Entries) - 1] +loop: + if entry == nil { + return + } + if !f(entry.Key, entry.Value) { + return + } + // Find current entry position in current node + e, _ := tree.search(node, entry.Key) + // Try to go down to the child left of the current entry + if e < len(node.Children) { + node = node.Children[e] + // Try to go down to the child right of the current node + for len(node.Children) > 0 { + node = node.Children[len(node.Children) - 1] + } + // Return the right-most entry + entry = node.Entries[len(node.Entries) - 1] + goto loop + } + // Above assures that we have reached a leaf node, so return the previous entry in current node (if any) + if e - 1 >= 0 { + entry = node.Entries[e - 1] + goto loop + } + + // Reached leaf node and there are no entries to the left of the current entry, so go up to the parent + for node.Parent != nil { + node = node.Parent + // Find previous entry position in current node (note: search returns the first equal or bigger than entry) + e, _ := tree.search(node, entry.Key) + // Check that there is a previous entry position in current node + if e - 1 >= 0 { + entry = node.Entries[e - 1] + goto loop + } + } +} + +func (tree *BTree) output(buffer *bytes.Buffer, node *BTreeNode, level int, isTail bool) { + for e := 0; e < len(node.Entries)+1; e++ { + if e < len(node.Children) { + tree.output(buffer, node.Children[e], level+1, true) + } + if e < len(node.Entries) { + if _, err := buffer.WriteString(strings.Repeat(" ", level)); err != nil { + } + if _, err := buffer.WriteString(fmt.Sprintf("%v", node.Entries[e].Key) + "\n"); err != nil { + } + } + } +} + +func (node *BTreeNode) height() int { + h := 0 + n := node + for ; n != nil; n = n.Children[0] { + h++ + if len(n.Children) == 0 { + break + } + } + return h +} + +func (tree *BTree) isLeaf(node *BTreeNode) bool { + return len(node.Children) == 0 +} + +func (tree *BTree) isFull(node *BTreeNode) bool { + return len(node.Entries) == tree.maxEntries() +} + +func (tree *BTree) shouldSplit(node *BTreeNode) bool { + return len(node.Entries) > tree.maxEntries() +} + +func (tree *BTree) maxChildren() int { + return tree.m +} + +func (tree *BTree) minChildren() int { + return (tree.m + 1) / 2 // ceil(m/2) +} + +func (tree *BTree) maxEntries() int { + return tree.maxChildren() - 1 +} + +func (tree *BTree) minEntries() int { + return tree.minChildren() - 1 +} + +func (tree *BTree) middle() int { + // "-1" to favor right nodes to have more keys when splitting + return (tree.m - 1) / 2 +} + +// search searches only within the single node among its entries +func (tree *BTree) search(node *BTreeNode, key interface{}) (index int, found bool) { + low, mid, high := 0, 0, len(node.Entries) - 1 + for low <= high { + mid = (high + low) / 2 + compare := tree.comparator(key, node.Entries[mid].Key) + switch { + case compare > 0: low = mid + 1 + case compare < 0: high = mid - 1 + case compare == 0: return mid, true + } + } + return low, false +} + +// searchRecursively searches recursively down the tree starting at the startNode +func (tree *BTree) searchRecursively(startNode *BTreeNode, key interface{}) (node *BTreeNode, index int, found bool) { + if tree.IsEmpty() { + return nil, -1, false + } + node = startNode + for { + index, found = tree.search(node, key) + if found { + return node, index, true + } + if tree.isLeaf(node) { + return nil, -1, false + } + node = node.Children[index] + } +} + +func (tree *BTree) insert(node *BTreeNode, entry *BTreeNodeEntry) (inserted bool) { + if tree.isLeaf(node) { + return tree.insertIntoLeaf(node, entry) + } + return tree.insertIntoInternal(node, entry) +} + +func (tree *BTree) insertIntoLeaf(node *BTreeNode, entry *BTreeNodeEntry) (inserted bool) { + insertPosition, found := tree.search(node, entry.Key) + if found { + node.Entries[insertPosition] = entry + return false + } + // Insert entry's key in the middle of the node + node.Entries = append(node.Entries, nil) + copy(node.Entries[insertPosition+1:], node.Entries[insertPosition:]) + node.Entries[insertPosition] = entry + tree.split(node) + return true +} + +func (tree *BTree) insertIntoInternal(node *BTreeNode, entry *BTreeNodeEntry) (inserted bool) { + insertPosition, found := tree.search(node, entry.Key) + if found { + node.Entries[insertPosition] = entry + return false + } + return tree.insert(node.Children[insertPosition], entry) +} + +func (tree *BTree) split(node *BTreeNode) { + if !tree.shouldSplit(node) { + return + } + + if node == tree.root { + tree.splitRoot() + return + } + + tree.splitNonRoot(node) +} + +func (tree *BTree) splitNonRoot(node *BTreeNode) { + middle := tree.middle() + parent := node.Parent + + left := &BTreeNode{Entries: append([]*BTreeNodeEntry(nil), node.Entries[:middle]...), Parent: parent} + right := &BTreeNode{Entries: append([]*BTreeNodeEntry(nil), node.Entries[middle+1:]...), Parent: parent} + + // Move children from the node to be split into left and right nodes + if !tree.isLeaf(node) { + left.Children = append([]*BTreeNode(nil), node.Children[:middle+1]...) + right.Children = append([]*BTreeNode(nil), node.Children[middle+1:]...) + setParent(left.Children, left) + setParent(right.Children, right) + } + + insertPosition, _ := tree.search(parent, node.Entries[middle].Key) + + // Insert middle key into parent + parent.Entries = append(parent.Entries, nil) + copy(parent.Entries[insertPosition+1:], parent.Entries[insertPosition:]) + parent.Entries[insertPosition] = node.Entries[middle] + + // Set child left of inserted key in parent to the created left node + parent.Children[insertPosition] = left + + // Set child right of inserted key in parent to the created right node + parent.Children = append(parent.Children, nil) + copy(parent.Children[insertPosition+2:], parent.Children[insertPosition+1:]) + parent.Children[insertPosition+1] = right + + tree.split(parent) +} + +func (tree *BTree) splitRoot() { + middle := tree.middle() + left := &BTreeNode{Entries: append([]*BTreeNodeEntry(nil), tree.root.Entries[:middle]...)} + right := &BTreeNode{Entries: append([]*BTreeNodeEntry(nil), tree.root.Entries[middle+1:]...)} + + // Move children from the node to be split into left and right nodes + if !tree.isLeaf(tree.root) { + left.Children = append([]*BTreeNode(nil), tree.root.Children[:middle+1]...) + right.Children = append([]*BTreeNode(nil), tree.root.Children[middle+1:]...) + setParent(left.Children, left) + setParent(right.Children, right) + } + + // Root is a node with one entry and two children (left and right) + newRoot := &BTreeNode{ + Entries: []*BTreeNodeEntry{tree.root.Entries[middle]}, + Children: []*BTreeNode{left, right}, + } + + left.Parent = newRoot + right.Parent = newRoot + tree.root = newRoot +} + +func setParent(nodes []*BTreeNode, parent *BTreeNode) { + for _, node := range nodes { + node.Parent = parent + } +} + +func (tree *BTree) left(node *BTreeNode) *BTreeNode { + if tree.IsEmpty() { + return nil + } + current := node + for { + if tree.isLeaf(current) { + return current + } + current = current.Children[0] + } +} + +func (tree *BTree) right(node *BTreeNode) *BTreeNode { + if tree.IsEmpty() { + return nil + } + current := node + for { + if tree.isLeaf(current) { + return current + } + current = current.Children[len(current.Children)-1] + } +} + +// leftSibling returns the node's left sibling and child index (in parent) if it exists, otherwise (nil,-1) +// key is any of keys in node (could even be deleted). +func (tree *BTree) leftSibling(node *BTreeNode, key interface{}) (*BTreeNode, int) { + if node.Parent != nil { + index, _ := tree.search(node.Parent, key) + index-- + if index >= 0 && index < len(node.Parent.Children) { + return node.Parent.Children[index], index + } + } + return nil, -1 +} + +// rightSibling returns the node's right sibling and child index (in parent) if it exists, otherwise (nil,-1) +// key is any of keys in node (could even be deleted). +func (tree *BTree) rightSibling(node *BTreeNode, key interface{}) (*BTreeNode, int) { + if node.Parent != nil { + index, _ := tree.search(node.Parent, key) + index++ + if index < len(node.Parent.Children) { + return node.Parent.Children[index], index + } + } + return nil, -1 +} + +// delete deletes an entry in node at entries' index +// ref.: https://en.wikipedia.org/wiki/B-tree#Deletion +func (tree *BTree) delete(node *BTreeNode, index int) { + // deleting from a leaf node + if tree.isLeaf(node) { + deletedKey := node.Entries[index].Key + tree.deleteEntry(node, index) + tree.rebalance(node, deletedKey) + if len(tree.root.Entries) == 0 { + tree.root = nil + } + return + } + + // deleting from an internal node + leftLargestNode := tree.right(node.Children[index]) // largest node in the left sub-tree (assumed to exist) + leftLargestEntryIndex := len(leftLargestNode.Entries) - 1 + node.Entries[index] = leftLargestNode.Entries[leftLargestEntryIndex] + deletedKey := leftLargestNode.Entries[leftLargestEntryIndex].Key + tree.deleteEntry(leftLargestNode, leftLargestEntryIndex) + tree.rebalance(leftLargestNode, deletedKey) +} + +// rebalance rebalances the tree after deletion if necessary and returns true, otherwise false. +// Note that we first delete the entry and then call rebalance, thus the passed deleted key as reference. +func (tree *BTree) rebalance(node *BTreeNode, deletedKey interface{}) { + // check if rebalancing is needed + if node == nil || len(node.Entries) >= tree.minEntries() { + return + } + + // try to borrow from left sibling + leftSibling, leftSiblingIndex := tree.leftSibling(node, deletedKey) + if leftSibling != nil && len(leftSibling.Entries) > tree.minEntries() { + // rotate right + node.Entries = append([]*BTreeNodeEntry{node.Parent.Entries[leftSiblingIndex]}, node.Entries...) // prepend parent's separator entry to node's entries + node.Parent.Entries[leftSiblingIndex] = leftSibling.Entries[len(leftSibling.Entries)-1] + tree.deleteEntry(leftSibling, len(leftSibling.Entries)-1) + if !tree.isLeaf(leftSibling) { + leftSiblingRightMostChild := leftSibling.Children[len(leftSibling.Children)-1] + leftSiblingRightMostChild.Parent = node + node.Children = append([]*BTreeNode{leftSiblingRightMostChild}, node.Children...) + tree.deleteChild(leftSibling, len(leftSibling.Children)-1) + } + return + } + + // try to borrow from right sibling + rightSibling, rightSiblingIndex := tree.rightSibling(node, deletedKey) + if rightSibling != nil && len(rightSibling.Entries) > tree.minEntries() { + // rotate left + node.Entries = append(node.Entries, node.Parent.Entries[rightSiblingIndex-1]) // append parent's separator entry to node's entries + node.Parent.Entries[rightSiblingIndex-1] = rightSibling.Entries[0] + tree.deleteEntry(rightSibling, 0) + if !tree.isLeaf(rightSibling) { + rightSiblingLeftMostChild := rightSibling.Children[0] + rightSiblingLeftMostChild.Parent = node + node.Children = append(node.Children, rightSiblingLeftMostChild) + tree.deleteChild(rightSibling, 0) + } + return + } + + // merge with siblings + if rightSibling != nil { + // merge with right sibling + node.Entries = append(node.Entries, node.Parent.Entries[rightSiblingIndex-1]) + node.Entries = append(node.Entries, rightSibling.Entries...) + deletedKey = node.Parent.Entries[rightSiblingIndex-1].Key + tree.deleteEntry(node.Parent, rightSiblingIndex-1) + tree.appendChildren(node.Parent.Children[rightSiblingIndex], node) + tree.deleteChild(node.Parent, rightSiblingIndex) + } else if leftSibling != nil { + // merge with left sibling + entries := append([]*BTreeNodeEntry(nil), leftSibling.Entries...) + entries = append(entries, node.Parent.Entries[leftSiblingIndex]) + node.Entries = append(entries, node.Entries...) + deletedKey = node.Parent.Entries[leftSiblingIndex].Key + tree.deleteEntry(node.Parent, leftSiblingIndex) + tree.prependChildren(node.Parent.Children[leftSiblingIndex], node) + tree.deleteChild(node.Parent, leftSiblingIndex) + } + + // make the merged node the root if its parent was the root and the root is empty + if node.Parent == tree.root && len(tree.root.Entries) == 0 { + tree.root = node + node.Parent = nil + return + } + + // parent might underflow, so try to rebalance if necessary + tree.rebalance(node.Parent, deletedKey) +} + +func (tree *BTree) prependChildren(fromNode *BTreeNode, toNode *BTreeNode) { + children := append([]*BTreeNode(nil), fromNode.Children...) + toNode.Children = append(children, toNode.Children...) + setParent(fromNode.Children, toNode) +} + +func (tree *BTree) appendChildren(fromNode *BTreeNode, toNode *BTreeNode) { + toNode.Children = append(toNode.Children, fromNode.Children...) + setParent(fromNode.Children, toNode) +} + +func (tree *BTree) deleteEntry(node *BTreeNode, index int) { + copy(node.Entries[index:], node.Entries[index+1:]) + node.Entries[len(node.Entries)-1] = nil + node.Entries = node.Entries[:len(node.Entries)-1] +} + +func (tree *BTree) deleteChild(node *BTreeNode, index int) { + if index >= len(node.Children) { + return + } + copy(node.Children[index:], node.Children[index+1:]) + node.Children[len(node.Children)-1] = nil + node.Children = node.Children[:len(node.Children)-1] +} \ No newline at end of file diff --git a/g/container/gtree/gtree_redblacktree.go b/g/container/gtree/gtree_redblacktree.go index 38737915a..8512c0e86 100644 --- a/g/container/gtree/gtree_redblacktree.go +++ b/g/container/gtree/gtree_redblacktree.go @@ -18,7 +18,7 @@ const ( black, red color = true, false ) -// RedBlackTree holds elements of the red-black tree +// RedBlackTree holds elements of the red-black tree. type RedBlackTree struct { mu *rwmutex.RWMutex root *RedBlackTreeNode @@ -26,17 +26,17 @@ type RedBlackTree struct { comparator func(v1, v2 interface{}) int } -// RedBlackTreeNode is a single element within the tree +// RedBlackTreeNode is a single element within the tree. type RedBlackTreeNode struct { - key interface{} - value interface{} + Key interface{} + Value interface{} color color left *RedBlackTreeNode right *RedBlackTreeNode parent *RedBlackTreeNode } -// NewWith instantiates a red-black tree with the custom comparator. +// NewRedBlackTree instantiates a red-black tree with the custom comparator. // The param used to specify whether using tree in un-concurrent-safety, // which is false in default. func NewRedBlackTree(comparator func(v1, v2 interface{}) int, unsafe...bool) *RedBlackTree { @@ -76,21 +76,21 @@ func (tree *RedBlackTree) doSet(key interface{}, value interface{}) { if tree.root == nil { // Assert key is of comparator's type for initial tree tree.comparator(key, key) - tree.root = &RedBlackTreeNode{key: key, value: value, color: red} + tree.root = &RedBlackTreeNode{Key: key, Value: value, color: red} insertedNode = tree.root } else { node := tree.root loop := true for loop { - compare := tree.comparator(key, node.key) + compare := tree.comparator(key, node.Key) switch { case compare == 0: - //node.key = key - node.value = value + //node.Key = key + node.Value = value return case compare < 0: if node.left == nil { - node.left = &RedBlackTreeNode{key: key, value: value, color: red} + node.left = &RedBlackTreeNode{Key: key, Value: value, color: red} insertedNode = node.left loop = false } else { @@ -98,7 +98,7 @@ func (tree *RedBlackTree) doSet(key interface{}, value interface{}) { } case compare > 0: if node.right == nil { - node.right = &RedBlackTreeNode{key: key, value: value, color: red} + node.right = &RedBlackTreeNode{Key: key, Value: value, color: red} insertedNode = node.right loop = false } else { @@ -131,7 +131,7 @@ func (tree *RedBlackTree) doSetWithLockCheck(key interface{}, value interface{}) tree.mu.Lock() defer tree.mu.Unlock() if node := tree.doSearch(key); node != nil { - return node.value + return node.Value } if f, ok := value.(func() interface {}); ok { value = f() @@ -247,8 +247,8 @@ func (tree *RedBlackTree) doRemove(key interface{}) (value interface{}) { } if node.left != nil && node.right != nil { p := node.left.maximumNode() - node.key = p.key - node.value = p.value + node.Key = p.Key + node.Value = p.Value node = p } if node.left == nil || node.right == nil { @@ -267,7 +267,7 @@ func (tree *RedBlackTree) doRemove(key interface{}) (value interface{}) { } } tree.size-- - value = node.value + value = node.Value return } @@ -337,6 +337,32 @@ func (tree *RedBlackTree) Map() map[interface{}]interface{} { func (tree *RedBlackTree) Left() *RedBlackTreeNode { tree.mu.RLock() defer tree.mu.RUnlock() + node := tree.leftNode() + if tree.mu.IsSafe() { + return &RedBlackTreeNode{ + Key : node.Key, + Value : node.Value, + } + } + return node +} + +// Right returns the right-most (max) node or nil if tree is empty. +func (tree *RedBlackTree) Right() *RedBlackTreeNode { + tree.mu.RLock() + defer tree.mu.RUnlock() + node := tree.rightNode() + if tree.mu.IsSafe() { + return &RedBlackTreeNode{ + Key : node.Key, + Value : node.Value, + } + } + return node +} + +// leftNode returns the left-most (min) node or nil if tree is empty. +func (tree *RedBlackTree) leftNode() *RedBlackTreeNode { p := (*RedBlackTreeNode)(nil) n := tree.root for n != nil { @@ -346,10 +372,8 @@ func (tree *RedBlackTree) Left() *RedBlackTreeNode { return p } -// Right returns the right-most (max) node or nil if tree is empty. -func (tree *RedBlackTree) Right() *RedBlackTreeNode { - tree.mu.RLock() - defer tree.mu.RUnlock() +// rightNode returns the right-most (max) node or nil if tree is empty. +func (tree *RedBlackTree) rightNode() *RedBlackTreeNode { p := (*RedBlackTreeNode)(nil) n := tree.root for n != nil { @@ -370,7 +394,7 @@ func (tree *RedBlackTree) Floor(key interface{}) (floor *RedBlackTreeNode) { found := false node := tree.root for node != nil { - compare := tree.comparator(key, node.key) + compare := tree.comparator(key, node.Key) switch { case compare == 0: return node @@ -398,7 +422,7 @@ func (tree *RedBlackTree) Ceiling(key interface{}) (ceiling *RedBlackTreeNode) { found := false node := tree.root for node != nil { - compare := tree.comparator(key, node.key) + compare := tree.comparator(key, node.Key) switch { case compare == 0: return node @@ -415,77 +439,71 @@ func (tree *RedBlackTree) Ceiling(key interface{}) (ceiling *RedBlackTreeNode) { return nil } -// IteratorAsc iterates the tree in ascendent order with given callback function . +// IteratorAsc iterates the tree in ascending order with given callback function . // If returns true, then it continues iterating; or false to stop. func (tree *RedBlackTree) IteratorAsc(f func (key, value interface{}) bool) { tree.mu.RLock() defer tree.mu.RUnlock() - node := tree.Left() + node := tree.leftNode() if node == nil { return } - for { - loop: - if node == nil { - break +loop: + if node == nil { + return + } + if !f(node.Key, node.Value) { + return + } + if node.right != nil { + node = node.right + for node.left != nil { + node = node.left } - if !f(node.key, node.value) { - break - } - if node.right != nil { - node = node.right - for node.left != nil { - node = node.left - } - continue - } - if node.parent != nil { - old := node - for node.parent != nil { - node = node.parent - if tree.comparator(old.key, node.key) <= 0 { - goto loop - } + goto loop + } + if node.parent != nil { + old := node + for node.parent != nil { + node = node.parent + if tree.comparator(old.Key, node.Key) <= 0 { + goto loop } } - break } } -// IteratorDesc iterates the tree in descendent order with given callback function . +// IteratorDesc iterates the tree in descending order with given callback function . // If returns true, then it continues iterating; or false to stop. func (tree *RedBlackTree) IteratorDesc(f func (key, value interface{}) bool) { tree.mu.RLock() defer tree.mu.RUnlock() - node := tree.Right() + node := tree.rightNode() if node == nil { return } - for { - loop: - if node == nil { - break +loop: + if node == nil { + return + } + if !f(node.Key, node.Value) { + return + } + if node.left != nil { + node = node.left + for node.right != nil { + node = node.right } - if !f(node.key, node.value) { - break - } - if node.left != nil { - node = node.left - for node.right != nil { - node = node.right - } - continue - } - if node.parent != nil { - old := node - for node.parent != nil { - node = node.parent - if tree.comparator(old.key, node.key) >= 0 { - goto loop - } + goto loop + } + if node.parent != nil { + old := node + for node.parent != nil { + node = node.parent + if tree.comparator(old.Key, node.Key) >= 0 { + goto loop } } - break } } @@ -501,7 +519,7 @@ func (tree *RedBlackTree) Clear() { func (tree *RedBlackTree) String() string { tree.mu.RLock() defer tree.mu.RUnlock() - str := "ROOT\n" + str := "RedBlackTree\n" if tree.size != 0 { tree.output(tree.root, "", true, &str) } @@ -514,7 +532,7 @@ func (tree *RedBlackTree) Print() { } func (node *RedBlackTreeNode) String() string { - return fmt.Sprintf("%v", node.key) + return fmt.Sprintf("%v", node.Key) } // Search searches the tree with given . @@ -524,7 +542,7 @@ func (tree *RedBlackTree) Search(key interface{}) (value interface{}, found bool defer tree.mu.RUnlock() node := tree.doSearch(key) if node != nil { - return node.value, true + return node.Value, true } return nil, false } @@ -584,7 +602,7 @@ func (tree *RedBlackTree) output(node *RedBlackTreeNode, prefix string, isTail b func (tree *RedBlackTree) doSearch(key interface{}) *RedBlackTreeNode { node := tree.root for node != nil { - compare := tree.comparator(key, node.key) + compare := tree.comparator(key, node.Key) switch { case compare == 0: return node case compare < 0: node = node.left diff --git a/geg/container/gtree/gtree_btree.go b/geg/container/gtree/gtree_btree.go new file mode 100644 index 000000000..01c001f4d --- /dev/null +++ b/geg/container/gtree/gtree_btree.go @@ -0,0 +1,21 @@ +package main + +import ( + "fmt" + "github.com/gogf/gf/g/container/gtree" +) + +func main() { + tree := gtree.NewBTree(10, func(v1, v2 interface{}) int { + return v1.(int) - v2.(int) + }) + for i := 0; i < 20; i++ { + tree.Set(i, i*10) + } + fmt.Println(tree.String()) + + tree.IteratorDesc(func(key, value interface{}) bool { + fmt.Println(key, value) + return true + }) +} \ No newline at end of file