diff --git a/container/gtree/gtree_redblacktree.go b/container/gtree/gtree_redblacktree.go index e9da6a07f6c..429be2fb6fe 100644 --- a/container/gtree/gtree_redblacktree.go +++ b/container/gtree/gtree_redblacktree.go @@ -7,38 +7,30 @@ package gtree import ( - "bytes" "fmt" + "github.com/emirpasic/gods/trees/redblacktree" "github.com/gogf/gf/v2/container/gvar" "github.com/gogf/gf/v2/internal/json" "github.com/gogf/gf/v2/internal/rwmutex" + "github.com/gogf/gf/v2/text/gstr" "github.com/gogf/gf/v2/util/gconv" "github.com/gogf/gf/v2/util/gutil" ) -type color bool - -const ( - black, red color = true, false -) +var _ iTree = (*RedBlackTree)(nil) // RedBlackTree holds elements of the red-black tree. type RedBlackTree struct { mu rwmutex.RWMutex - root *RedBlackTreeNode - size int comparator func(v1, v2 interface{}) int + tree *redblacktree.Tree } // RedBlackTreeNode is a single element within the tree. type RedBlackTreeNode struct { - Key interface{} - Value interface{} - color color - left *RedBlackTreeNode - right *RedBlackTreeNode - parent *RedBlackTreeNode + Key interface{} + Value interface{} } // NewRedBlackTree instantiates a red-black tree with the custom key comparator. @@ -48,6 +40,7 @@ func NewRedBlackTree(comparator func(v1, v2 interface{}) int, safe ...bool) *Red return &RedBlackTree{ mu: rwmutex.Create(safe...), comparator: comparator, + tree: redblacktree.NewWith(comparator), } } @@ -64,21 +57,14 @@ func NewRedBlackTreeFrom(comparator func(v1, v2 interface{}) int, data map[inter // SetComparator sets/changes the comparator for sorting. func (tree *RedBlackTree) SetComparator(comparator func(a, b interface{}) int) { - tree.mu.Lock() - defer tree.mu.Unlock() tree.comparator = comparator - if tree.size > 0 { - data := make(map[interface{}]interface{}, tree.size) - tree.doIteratorAsc(tree.leftNode(), func(key, value interface{}) bool { - data[key] = value - return true - }) - // Resort the tree if comparator is changed. - tree.root = nil - tree.size = 0 - for k, v := range data { - tree.doSet(k, v) - } + if tree.tree == nil { + tree.tree = redblacktree.NewWith(comparator) + } + size := tree.tree.Size() + if size > 0 { + m := tree.Map() + tree.Sets(m) } } @@ -89,7 +75,7 @@ func (tree *RedBlackTree) Clone() *RedBlackTree { return newTree } -// Set inserts key-value item into the tree. +// Set inserts node into the tree. func (tree *RedBlackTree) Set(key interface{}, value interface{}) { tree.mu.Lock() defer tree.mu.Unlock() @@ -100,88 +86,63 @@ func (tree *RedBlackTree) Set(key interface{}, value interface{}) { func (tree *RedBlackTree) Sets(data map[interface{}]interface{}) { tree.mu.Lock() defer tree.mu.Unlock() - for k, v := range data { - tree.doSet(k, v) + for key, value := range data { + tree.doSet(key, value) } } -// doSet inserts key-value item into the tree without mutex. -func (tree *RedBlackTree) doSet(key interface{}, value interface{}) { - insertedNode := (*RedBlackTreeNode)(nil) - if tree.root == nil { - // Assert key is of comparator's type for initial tree - tree.getComparator()(key, key) - tree.root = &RedBlackTreeNode{Key: key, Value: value, color: red} - insertedNode = tree.root - } else { - node := tree.root - loop := true - for loop { - compare := tree.getComparator()(key, node.Key) - switch { - case compare == 0: - // node.Key = key - node.Value = value - return - case compare < 0: - if node.left == nil { - node.left = &RedBlackTreeNode{Key: key, Value: value, color: red} - insertedNode = node.left - loop = false - } else { - node = node.left - } - case compare > 0: - if node.right == nil { - node.right = &RedBlackTreeNode{Key: key, Value: value, color: red} - insertedNode = node.right - loop = false - } else { - node = node.right - } - } - } - insertedNode.parent = node +// SetIfNotExist sets `value` to the map if the `key` does not exist, and then returns true. +// It returns false if `key` exists, and `value` would be ignored. +func (tree *RedBlackTree) SetIfNotExist(key interface{}, value interface{}) bool { + tree.mu.Lock() + defer tree.mu.Unlock() + if _, ok := tree.doGet(key); !ok { + tree.doSet(key, value) + return true } - tree.insertCase1(insertedNode) - tree.size++ + return false } -// Get searches the node in the tree by `key` and returns its value or nil if key is not found in tree. -func (tree *RedBlackTree) Get(key interface{}) (value interface{}) { - value, _ = tree.Search(key) - return +// SetIfNotExistFunc sets value with return value of callback function `f`, and then returns true. +// It returns false if `key` exists, and `value` would be ignored. +func (tree *RedBlackTree) SetIfNotExistFunc(key interface{}, f func() interface{}) bool { + tree.mu.Lock() + defer tree.mu.Unlock() + if _, ok := tree.doGet(key); !ok { + tree.doSet(key, f()) + return true + } + return false } -// 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 , -// it will be executed with mutex.Lock of the hash map, -// and its return value will be set to the map with `key`. +// SetIfNotExistFuncLock sets value with return value of callback function `f`, and then returns true. +// It returns false if `key` exists, and `value` would be ignored. // -// It returns value with given `key`. -func (tree *RedBlackTree) doSetWithLockCheck(key interface{}, value interface{}) interface{} { +// SetIfNotExistFuncLock differs with SetIfNotExistFunc function is that +// it executes function `f` with mutex.Lock of the hash map. +func (tree *RedBlackTree) SetIfNotExistFuncLock(key interface{}, f func() interface{}) bool { tree.mu.Lock() defer tree.mu.Unlock() - if node, found := tree.doSearch(key); found { - return node.Value - } - if f, ok := value.(func() interface{}); ok { - value = f() - } - if value != nil { - tree.doSet(key, value) + if _, ok := tree.doGet(key); !ok { + tree.doSet(key, f) + return true } - return value + return false +} + +// Get searches the node in the tree by `key` and returns its value or nil if key is not found in tree. +func (tree *RedBlackTree) Get(key interface{}) (value interface{}) { + value, _ = tree.Search(key) + return } // GetOrSet returns the value by key, // or sets value with given `value` if it does not exist and then returns this value. func (tree *RedBlackTree) GetOrSet(key interface{}, value interface{}) interface{} { - if v, ok := tree.Search(key); !ok { - return tree.doSetWithLockCheck(key, value) + tree.mu.Lock() + defer tree.mu.Unlock() + if v, ok := tree.doGet(key); !ok { + return tree.doSet(key, value) } else { return v } @@ -191,8 +152,10 @@ func (tree *RedBlackTree) GetOrSet(key interface{}, value interface{}) interface // or sets value with returned value of callback function `f` if it does not exist // and then returns this value. func (tree *RedBlackTree) GetOrSetFunc(key interface{}, f func() interface{}) interface{} { - if v, ok := tree.Search(key); !ok { - return tree.doSetWithLockCheck(key, f()) + tree.mu.Lock() + defer tree.mu.Unlock() + if v, ok := tree.doGet(key); !ok { + return tree.doSet(key, f()) } else { return v } @@ -205,8 +168,10 @@ func (tree *RedBlackTree) GetOrSetFunc(key interface{}, f func() interface{}) in // GetOrSetFuncLock differs with GetOrSetFunc function is that it executes function `f` // with mutex.Lock of the hash map. func (tree *RedBlackTree) GetOrSetFuncLock(key interface{}, f func() interface{}) interface{} { - if v, ok := tree.Search(key); !ok { - return tree.doSetWithLockCheck(key, f) + tree.mu.Lock() + defer tree.mu.Unlock() + if v, ok := tree.doGet(key); !ok { + return tree.doSet(key, f) } else { return v } @@ -236,79 +201,41 @@ func (tree *RedBlackTree) GetVarOrSetFuncLock(key interface{}, f func() interfac return gvar.New(tree.GetOrSetFuncLock(key, f)) } -// SetIfNotExist sets `value` to the map if the `key` does not exist, and then returns true. -// It returns false if `key` exists, and `value` would be ignored. -func (tree *RedBlackTree) 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`, and then returns true. -// It returns false if `key` exists, and `value` would be ignored. -func (tree *RedBlackTree) 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`, and then returns 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 *RedBlackTree) SetIfNotExistFuncLock(key interface{}, f func() interface{}) bool { - if !tree.Contains(key) { - tree.doSetWithLockCheck(key, f) - return true +// Search searches the tree with given `key`. +// Second return parameter `found` is true if key was found, otherwise false. +func (tree *RedBlackTree) Search(key interface{}) (value interface{}, found bool) { + tree.mu.RLock() + defer tree.mu.RUnlock() + if node, found := tree.doGet(key); found { + return node, true } - return false + return nil, false } // Contains checks whether `key` exists in the tree. func (tree *RedBlackTree) Contains(key interface{}) bool { - _, ok := tree.Search(key) + tree.mu.RLock() + defer tree.mu.RUnlock() + _, ok := tree.doGet(key) return ok } -// doRemove removes the node from the tree by `key` without mutex. -func (tree *RedBlackTree) doRemove(key interface{}) (value interface{}) { - child := (*RedBlackTreeNode)(nil) - node, found := tree.doSearch(key) - if !found { - return - } - value = node.Value - if node.left != nil && node.right != nil { - p := node.left.maximumNode() - node.Key = p.Key - node.Value = p.Value - node = p - } - if node.left == nil || node.right == nil { - if node.right == nil { - child = node.left - } else { - child = node.right - } - if node.color == black { - node.color = tree.nodeColor(child) - tree.deleteCase1(node) - } - tree.replaceNode(node, child) - if node.parent == nil && child != nil { - child.color = black - } - } - tree.size-- - return +// Size returns number of nodes in the tree. +func (tree *RedBlackTree) Size() int { + tree.mu.RLock() + defer tree.mu.RUnlock() + return tree.tree.Size() } -// Remove removes the node from the tree by `key`. +// IsEmpty returns true if tree does not contain any nodes. +func (tree *RedBlackTree) IsEmpty() bool { + tree.mu.RLock() + defer tree.mu.RUnlock() + return tree.tree.Size() == 0 +} + +// Remove removes the node from the tree by key. +// Key should adhere to the comparator's type assertion, otherwise method panics. func (tree *RedBlackTree) Remove(key interface{}) (value interface{}) { tree.mu.Lock() defer tree.mu.Unlock() @@ -324,170 +251,78 @@ func (tree *RedBlackTree) Removes(keys []interface{}) { } } -// IsEmpty returns true if tree does not contain any nodes. -func (tree *RedBlackTree) IsEmpty() bool { - return tree.Size() == 0 -} - -// Size returns number of nodes in the tree. -func (tree *RedBlackTree) Size() int { - tree.mu.RLock() - defer tree.mu.RUnlock() - return tree.size +// Clear removes all nodes from the tree. +func (tree *RedBlackTree) Clear() { + tree.mu.Lock() + defer tree.mu.Unlock() + tree.tree.Clear() } // Keys returns all keys in asc order. func (tree *RedBlackTree) Keys() []interface{} { - var ( - keys = make([]interface{}, tree.Size()) - index = 0 - ) - tree.IteratorAsc(func(key, value interface{}) bool { - keys[index] = key - index++ - return true - }) - return keys + tree.mu.RLock() + defer tree.mu.RUnlock() + return tree.tree.Keys() } // Values returns all values in asc order based on the key. func (tree *RedBlackTree) Values() []interface{} { - var ( - values = make([]interface{}, tree.Size()) - index = 0 - ) - tree.IteratorAsc(func(key, value interface{}) bool { - values[index] = value - index++ - return true - }) - return values + tree.mu.RLock() + defer tree.mu.RUnlock() + return tree.tree.Values() } -// Map returns all key-value items as map. -func (tree *RedBlackTree) 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 +// Replace the data of the tree with given `data`. +func (tree *RedBlackTree) Replace(data map[interface{}]interface{}) { + tree.mu.Lock() + defer tree.mu.Unlock() + tree.tree.Clear() + for k, v := range data { + tree.doSet(k, v) + } } -// MapStrAny returns all key-value items as map[string]interface{}. -func (tree *RedBlackTree) MapStrAny() map[string]interface{} { - m := make(map[string]interface{}, tree.Size()) - tree.IteratorAsc(func(key, value interface{}) bool { - m[gconv.String(key)] = value - return true - }) - return m +// Print prints the tree to stdout. +func (tree *RedBlackTree) Print() { + fmt.Println(tree.String()) } -// Left returns the left-most (min) node or nil if tree is empty. -func (tree *RedBlackTree) Left() *RedBlackTreeNode { +// String returns a string representation of container +func (tree *RedBlackTree) String() string { tree.mu.RLock() defer tree.mu.RUnlock() - node := tree.leftNode() - if tree.mu.IsSafe() { - return &RedBlackTreeNode{ - Key: node.Key, - Value: node.Value, - } - } - return node + return gstr.Replace(tree.tree.String(), "RedBlackTree\n", "") } -// Right returns the right-most (max) node or nil if tree is empty. -func (tree *RedBlackTree) Right() *RedBlackTreeNode { +// MarshalJSON implements the interface MarshalJSON for json.Marshal. +func (tree *RedBlackTree) MarshalJSON() (jsonBytes []byte, err error) { 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 { - p = n - n = n.left - } - return p -} - -// 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 { - p = n - n = n.right - } - return p + return tree.tree.MarshalJSON() } -// Floor Finds floor node of the input key, return the floor node or nil if no floor node is found. -// Second return parameter is true if floor was found, otherwise false. -// -// Floor node is defined as the largest node that its key is smaller than or equal to the given `key`. -// A floor node may not be found, either because the tree is empty, or because -// all nodes in the tree are larger than the given node. -func (tree *RedBlackTree) Floor(key interface{}) (floor *RedBlackTreeNode, found bool) { +// Map returns all key-value items as map. +func (tree *RedBlackTree) Map() map[interface{}]interface{} { tree.mu.RLock() defer tree.mu.RUnlock() - n := tree.root - for n != nil { - compare := tree.getComparator()(key, n.Key) - switch { - case compare == 0: - return n, true - case compare < 0: - n = n.left - case compare > 0: - floor, found = n, true - n = n.right - } - } - if found { - return - } - return nil, false + m := make(map[interface{}]interface{}, tree.Size()) + tree.IteratorAsc(func(key, value interface{}) bool { + m[key] = value + return true + }) + return m } -// Ceiling finds ceiling node of the input key, return the ceiling node or nil if no ceiling node is found. -// Second return parameter is true if ceiling was found, otherwise false. -// -// Ceiling node is defined as the smallest node that its key is larger than or equal to the given `key`. -// A ceiling node may not be found, either because the tree is empty, or because -// all nodes in the tree are smaller than the given node. -func (tree *RedBlackTree) Ceiling(key interface{}) (ceiling *RedBlackTreeNode, found bool) { +// MapStrAny returns all key-value items as map[string]interface{}. +func (tree *RedBlackTree) MapStrAny() map[string]interface{} { tree.mu.RLock() defer tree.mu.RUnlock() - n := tree.root - for n != nil { - compare := tree.getComparator()(key, n.Key) - switch { - case compare == 0: - return n, true - case compare > 0: - n = n.right - case compare < 0: - ceiling, found = n, true - n = n.left - } - } - if found { - return - } - return nil, false + m := make(map[string]interface{}, tree.Size()) + tree.IteratorAsc(func(key, value interface{}) bool { + m[gconv.String(key)] = value + return true + }) + return m } // Iterator is alias of IteratorAsc. @@ -505,7 +340,13 @@ func (tree *RedBlackTree) IteratorFrom(key interface{}, match bool, f func(key, func (tree *RedBlackTree) IteratorAsc(f func(key, value interface{}) bool) { tree.mu.RLock() defer tree.mu.RUnlock() - tree.doIteratorAsc(tree.leftNode(), f) + it := tree.tree.Iterator() + for it.Begin(); it.Next(); { + index, value := it.Key(), it.Value() + if ok := f(index, value); !ok { + break + } + } } // IteratorAscFrom iterates the tree readonly in ascending order with given callback function `f`. @@ -515,39 +356,13 @@ func (tree *RedBlackTree) IteratorAsc(f func(key, value interface{}) bool) { func (tree *RedBlackTree) IteratorAscFrom(key interface{}, match bool, f func(key, value interface{}) bool) { tree.mu.RLock() defer tree.mu.RUnlock() - node, found := tree.doSearch(key) - if match { - if found { - tree.doIteratorAsc(node, f) - } - } else { - tree.doIteratorAsc(node, f) - } -} - -func (tree *RedBlackTree) doIteratorAsc(node *RedBlackTreeNode, f func(key, value interface{}) bool) { -loop: - if node == nil { + var keys = tree.tree.Keys() + index, isIterator := tree.iteratorFromGetIndex(key, keys, match) + if !isIterator { return } - if !f(node.Key, node.Value) { - return - } - if node.right != nil { - node = node.right - for node.left != nil { - node = node.left - } - goto loop - } - if node.parent != nil { - old := node - for node.parent != nil { - node = node.parent - if tree.getComparator()(old.Key, node.Key) <= 0 { - goto loop - } - } + for ; index < len(keys); index++ { + f(keys[index], tree.Get(keys[index])) } } @@ -556,7 +371,13 @@ loop: func (tree *RedBlackTree) IteratorDesc(f func(key, value interface{}) bool) { tree.mu.RLock() defer tree.mu.RUnlock() - tree.doIteratorDesc(tree.rightNode(), f) + it := tree.tree.Iterator() + for it.End(); it.Prev(); { + index, value := it.Key(), it.Value() + if ok := f(index, value); !ok { + break + } + } } // IteratorDescFrom iterates the tree readonly in descending order with given callback function `f`. @@ -566,90 +387,86 @@ func (tree *RedBlackTree) IteratorDesc(f func(key, value interface{}) bool) { func (tree *RedBlackTree) IteratorDescFrom(key interface{}, match bool, f func(key, value interface{}) bool) { tree.mu.RLock() defer tree.mu.RUnlock() - node, found := tree.doSearch(key) - if match { - if found { - tree.doIteratorDesc(node, f) - } - } else { - tree.doIteratorDesc(node, f) + var keys = tree.tree.Keys() + index, isIterator := tree.iteratorFromGetIndex(key, keys, match) + if !isIterator { + return + } + for ; index >= 0; index-- { + f(keys[index], tree.Get(keys[index])) } } -func (tree *RedBlackTree) doIteratorDesc(node *RedBlackTreeNode, f func(key, value interface{}) bool) { -loop: +// Left returns the minimum element of the AVL tree +// or nil if the tree is empty. +func (tree *RedBlackTree) Left() *RedBlackTreeNode { + tree.mu.RLock() + defer tree.mu.RUnlock() + node := tree.tree.Left() if node == nil { - return - } - if !f(node.Key, node.Value) { - return + return nil } - if node.left != nil { - node = node.left - for node.right != nil { - node = node.right - } - goto loop - } - if node.parent != nil { - old := node - for node.parent != nil { - node = node.parent - if tree.getComparator()(old.Key, node.Key) >= 0 { - goto loop - } - } + return &RedBlackTreeNode{ + Key: node.Key, + Value: node.Value, } } -// Clear removes all nodes from the tree. -func (tree *RedBlackTree) Clear() { - tree.mu.Lock() - defer tree.mu.Unlock() - tree.root = nil - tree.size = 0 -} - -// Replace the data of the tree with given `data`. -func (tree *RedBlackTree) Replace(data map[interface{}]interface{}) { - tree.mu.Lock() - defer tree.mu.Unlock() - tree.root = nil - tree.size = 0 - for k, v := range data { - tree.doSet(k, v) +// Right returns the maximum element of the AVL tree +// or nil if the tree is empty. +func (tree *RedBlackTree) Right() *RedBlackTreeNode { + tree.mu.RLock() + defer tree.mu.RUnlock() + node := tree.tree.Right() + if node == nil { + return nil + } + return &RedBlackTreeNode{ + Key: node.Key, + Value: node.Value, } } -// String returns a string representation of container. -func (tree *RedBlackTree) String() string { - if tree == nil { - return "" - } +// Floor Finds floor node of the input key, return the floor node or nil if no floor node is found. +// Second return parameter is true if floor was found, otherwise false. +// +// Floor node is defined as the largest node that is smaller than or equal to the given node. +// A floor node may not be found, either because the tree is empty, or because +// all nodes in the tree is larger than the given node. +// +// Key should adhere to the comparator's type assertion, otherwise method panics. +func (tree *RedBlackTree) Floor(key interface{}) (floor *RedBlackTreeNode, found bool) { tree.mu.RLock() defer tree.mu.RUnlock() - str := "" - if tree.size != 0 { - tree.output(tree.root, "", true, &str) + node, found := tree.tree.Floor(key) + if !found { + return nil, false } - return str + return &RedBlackTreeNode{ + Key: node.Key, + Value: node.Value, + }, true } -// Print prints the tree to stdout. -func (tree *RedBlackTree) Print() { - fmt.Println(tree.String()) -} - -// Search searches the tree with given `key`. -// Second return parameter `found` is true if key was found, otherwise false. -func (tree *RedBlackTree) Search(key interface{}) (value interface{}, found bool) { +// Ceiling finds ceiling node of the input key, return the ceiling node or nil if no ceiling node is found. +// Second return parameter is true if ceiling was found, otherwise false. +// +// Ceiling node is defined as the smallest node that is larger than or equal to the given node. +// A ceiling node may not be found, either because the tree is empty, or because +// all nodes in the tree is smaller than the given node. +// +// Key should adhere to the comparator's type assertion, otherwise method panics. +func (tree *RedBlackTree) Ceiling(key interface{}) (ceiling *RedBlackTreeNode, found bool) { tree.mu.RLock() defer tree.mu.RUnlock() - node, found := tree.doSearch(key) - if found { - return node.Value, true + node, found := tree.tree.Ceiling(key) + if !found { + return nil, false } - return nil, false + return &RedBlackTreeNode{ + Key: node.Key, + Value: node.Value, + }, true } // Flip exchanges key-value of the tree to value-key. @@ -658,7 +475,7 @@ func (tree *RedBlackTree) Search(key interface{}) (value interface{}, found bool // // If the type of value is different with key, you pass the new `comparator`. func (tree *RedBlackTree) Flip(comparator ...func(v1, v2 interface{}) int) { - t := (*RedBlackTree)(nil) + var t = new(RedBlackTree) if len(comparator) > 0 { t = NewRedBlackTree(comparator[0], tree.mu.IsSafe()) } else { @@ -668,287 +485,8 @@ func (tree *RedBlackTree) Flip(comparator ...func(v1, v2 interface{}) int) { t.doSet(value, key) return true }) - tree.mu.Lock() - tree.root = t.root - tree.size = t.size - tree.mu.Unlock() -} - -func (tree *RedBlackTree) output(node *RedBlackTreeNode, prefix string, isTail bool, str *string) { - if node.right != nil { - newPrefix := prefix - if isTail { - newPrefix += "│ " - } else { - newPrefix += " " - } - tree.output(node.right, newPrefix, false, str) - } - *str += prefix - if isTail { - *str += "└── " - } else { - *str += "┌── " - } - *str += fmt.Sprintf("%v\n", node.Key) - if node.left != nil { - newPrefix := prefix - if isTail { - newPrefix += " " - } else { - newPrefix += "│ " - } - tree.output(node.left, newPrefix, true, str) - } -} - -// doSearch searches the tree with given `key` without mutex. -// It returns the node if found or otherwise nil. -func (tree *RedBlackTree) doSearch(key interface{}) (node *RedBlackTreeNode, found bool) { - node = tree.root - for node != nil { - compare := tree.getComparator()(key, node.Key) - switch { - case compare == 0: - return node, true - case compare < 0: - node = node.left - case compare > 0: - node = node.right - } - } - return node, false -} - -func (node *RedBlackTreeNode) grandparent() *RedBlackTreeNode { - if node != nil && node.parent != nil { - return node.parent.parent - } - return nil -} - -func (node *RedBlackTreeNode) uncle() *RedBlackTreeNode { - if node == nil || node.parent == nil || node.parent.parent == nil { - return nil - } - return node.parent.sibling() -} - -func (node *RedBlackTreeNode) sibling() *RedBlackTreeNode { - if node == nil || node.parent == nil { - return nil - } - if node == node.parent.left { - return node.parent.right - } - return node.parent.left -} - -func (tree *RedBlackTree) rotateLeft(node *RedBlackTreeNode) { - right := node.right - tree.replaceNode(node, right) - node.right = right.left - if right.left != nil { - right.left.parent = node - } - right.left = node - node.parent = right -} - -func (tree *RedBlackTree) rotateRight(node *RedBlackTreeNode) { - left := node.left - tree.replaceNode(node, left) - node.left = left.right - if left.right != nil { - left.right.parent = node - } - left.right = node - node.parent = left -} - -func (tree *RedBlackTree) replaceNode(old *RedBlackTreeNode, new *RedBlackTreeNode) { - if old.parent == nil { - tree.root = new - } else { - if old == old.parent.left { - old.parent.left = new - } else { - old.parent.right = new - } - } - if new != nil { - new.parent = old.parent - } -} - -func (tree *RedBlackTree) insertCase1(node *RedBlackTreeNode) { - if node.parent == nil { - node.color = black - } else { - tree.insertCase2(node) - } -} - -func (tree *RedBlackTree) insertCase2(node *RedBlackTreeNode) { - if tree.nodeColor(node.parent) == black { - return - } - tree.insertCase3(node) -} - -func (tree *RedBlackTree) insertCase3(node *RedBlackTreeNode) { - uncle := node.uncle() - if tree.nodeColor(uncle) == red { - node.parent.color = black - uncle.color = black - node.grandparent().color = red - tree.insertCase1(node.grandparent()) - } else { - tree.insertCase4(node) - } -} - -func (tree *RedBlackTree) insertCase4(node *RedBlackTreeNode) { - grandparent := node.grandparent() - if node == node.parent.right && node.parent == grandparent.left { - tree.rotateLeft(node.parent) - node = node.left - } else if node == node.parent.left && node.parent == grandparent.right { - tree.rotateRight(node.parent) - node = node.right - } - tree.insertCase5(node) -} - -func (tree *RedBlackTree) insertCase5(node *RedBlackTreeNode) { - node.parent.color = black - grandparent := node.grandparent() - grandparent.color = red - if node == node.parent.left && node.parent == grandparent.left { - tree.rotateRight(grandparent) - } else if node == node.parent.right && node.parent == grandparent.right { - tree.rotateLeft(grandparent) - } -} - -func (node *RedBlackTreeNode) maximumNode() *RedBlackTreeNode { - if node == nil { - return nil - } - for node.right != nil { - return node.right - } - return node -} - -func (tree *RedBlackTree) deleteCase1(node *RedBlackTreeNode) { - if node.parent == nil { - return - } - tree.deleteCase2(node) -} - -func (tree *RedBlackTree) deleteCase2(node *RedBlackTreeNode) { - sibling := node.sibling() - if tree.nodeColor(sibling) == red { - node.parent.color = red - sibling.color = black - if node == node.parent.left { - tree.rotateLeft(node.parent) - } else { - tree.rotateRight(node.parent) - } - } - tree.deleteCase3(node) -} - -func (tree *RedBlackTree) deleteCase3(node *RedBlackTreeNode) { - sibling := node.sibling() - if tree.nodeColor(node.parent) == black && - tree.nodeColor(sibling) == black && - tree.nodeColor(sibling.left) == black && - tree.nodeColor(sibling.right) == black { - sibling.color = red - tree.deleteCase1(node.parent) - } else { - tree.deleteCase4(node) - } -} - -func (tree *RedBlackTree) deleteCase4(node *RedBlackTreeNode) { - sibling := node.sibling() - if tree.nodeColor(node.parent) == red && - tree.nodeColor(sibling) == black && - tree.nodeColor(sibling.left) == black && - tree.nodeColor(sibling.right) == black { - sibling.color = red - node.parent.color = black - } else { - tree.deleteCase5(node) - } -} - -func (tree *RedBlackTree) deleteCase5(node *RedBlackTreeNode) { - sibling := node.sibling() - if node == node.parent.left && - tree.nodeColor(sibling) == black && - tree.nodeColor(sibling.left) == red && - tree.nodeColor(sibling.right) == black { - sibling.color = red - sibling.left.color = black - tree.rotateRight(sibling) - } else if node == node.parent.right && - tree.nodeColor(sibling) == black && - tree.nodeColor(sibling.right) == red && - tree.nodeColor(sibling.left) == black { - sibling.color = red - sibling.right.color = black - tree.rotateLeft(sibling) - } - tree.deleteCase6(node) -} - -func (tree *RedBlackTree) deleteCase6(node *RedBlackTreeNode) { - sibling := node.sibling() - sibling.color = tree.nodeColor(node.parent) - node.parent.color = black - if node == node.parent.left && tree.nodeColor(sibling.right) == red { - sibling.right.color = black - tree.rotateLeft(node.parent) - } else if tree.nodeColor(sibling.left) == red { - sibling.left.color = black - tree.rotateRight(node.parent) - } -} - -func (tree *RedBlackTree) nodeColor(node *RedBlackTreeNode) color { - if node == nil { - return black - } - return node.color -} - -// MarshalJSON implements the interface MarshalJSON for json.Marshal. -func (tree RedBlackTree) MarshalJSON() (jsonBytes []byte, err error) { - if tree.root == nil { - return []byte("null"), nil - } - buffer := bytes.NewBuffer(nil) - buffer.WriteByte('{') - tree.Iterator(func(key, value interface{}) bool { - valueBytes, valueJsonErr := json.Marshal(value) - if valueJsonErr != nil { - err = valueJsonErr - return false - } - if buffer.Len() > 1 { - buffer.WriteByte(',') - } - buffer.WriteString(fmt.Sprintf(`"%v":%s`, key, valueBytes)) - return true - }) - buffer.WriteByte('}') - return buffer.Bytes(), nil + tree.Clear() + tree.Sets(t.Map()) } // UnmarshalJSON implements the interface UnmarshalJSON for json.Unmarshal. @@ -957,6 +495,7 @@ func (tree *RedBlackTree) UnmarshalJSON(b []byte) error { defer tree.mu.Unlock() if tree.comparator == nil { tree.comparator = gutil.ComparatorString + tree.tree = redblacktree.NewWith(tree.comparator) } var data map[string]interface{} if err := json.UnmarshalUseNumber(b, &data); err != nil { @@ -974,6 +513,7 @@ func (tree *RedBlackTree) UnmarshalValue(value interface{}) (err error) { defer tree.mu.Unlock() if tree.comparator == nil { tree.comparator = gutil.ComparatorString + tree.tree = redblacktree.NewWith(tree.comparator) } for k, v := range gconv.Map(value) { tree.doSet(k, v) @@ -981,11 +521,47 @@ func (tree *RedBlackTree) UnmarshalValue(value interface{}) (err error) { return } -// getComparator returns the comparator if it's previously set, -// or else it panics. -func (tree *RedBlackTree) getComparator() func(a, b interface{}) int { - if tree.comparator == nil { - panic("comparator is missing for tree") +// doSet sets key-value pair to the tree. +func (tree *RedBlackTree) doSet(key, value interface{}) interface{} { + if f, ok := value.(func() interface{}); ok { + value = f() + } + if value == nil { + return value + } + tree.tree.Put(key, value) + return value +} + +// doGet retrieves and returns the value of given key from tree. +func (tree *RedBlackTree) doGet(key interface{}) (value interface{}, found bool) { + return tree.tree.Get(key) +} + +// doRemove removes key from tree. +func (tree *RedBlackTree) doRemove(key interface{}) (value interface{}) { + value, _ = tree.tree.Get(key) + tree.tree.Remove(key) + return +} + +// iteratorFromGetIndex returns the index of the key in the keys slice. +// The parameter `match` specifies whether starting iterating if the `key` is fully matched, +// or else using index searching iterating. +// If `isIterator` is true, iterator is available; or else not. +func (tree *RedBlackTree) iteratorFromGetIndex(key interface{}, keys []interface{}, match bool) (index int, isIterator bool) { + if match { + for i, k := range keys { + if k == key { + isIterator = true + index = i + } + } + } else { + if i, ok := key.(int); ok { + isIterator = true + index = i + } } - return tree.comparator + return } diff --git a/container/gtree/gtree_z_example_redblacktree_test.go b/container/gtree/gtree_z_example_redblacktree_test.go index 552b6232daa..c28ab5d1d4b 100644 --- a/container/gtree/gtree_z_example_redblacktree_test.go +++ b/container/gtree/gtree_z_example_redblacktree_test.go @@ -592,6 +592,11 @@ func ExampleRedBlackTree_IteratorAscFrom_NoExistKeyAndMatchFalse() { }) // Output: + // key: 1 , value: 10 + // key: 2 , value: 20 + // key: 3 , value: 30 + // key: 4 , value: 40 + // key: 5 , value: 50 } func ExampleRedBlackTree_IteratorDesc() {