add BTree container for gtree package

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 <m> (maximum number of children) and a custom key comparator.
+// Note that the <m> 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 <key> 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 <key>,
+// or else just return the existing value.
+//
+// When setting value, if <value> is type of <func() interface {}>,
+// it will be executed with mutex.Lock of the hash map,
+// and its return value will be set to the map with <key>.
+//
+// It returns value with given <key>.
+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 <value> 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 <f> 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 <f> if not exist
+// and returns this value.
+//
+// GetOrSetFuncLock differs with GetOrSetFunc function is that it executes function <f>
+// 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 <key>.
+// 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 <value> to the map if the <key> does not exist, then return true.
+// It returns false if <key> exists, and <value> 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 <f>, then return true.
+// It returns false if <key> exists, and <value> 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 <f>, then return true.
+// It returns false if <key> exists, and <value> would be ignored.
+//
+// SetIfNotExistFuncLock differs with SetIfNotExistFunc function is that
+// it executes function <f> 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 <key> 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 <key>.
+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 <keys>.
+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 <key>.
+// Second return parameter <found> 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 <key> 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 <f>.
+// If <f> 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 <f>.
+// If <f> 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]
+}

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 <unsafe> 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 <f>.
+// IteratorAsc iterates the tree in ascending order with given callback function <f>.
 // If <f> 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 <f>.
+// IteratorDesc iterates the tree in descending order with given callback function <f>.
 // If <f> 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 <key>.
@@ -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

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
+	})
+}