diff --git a/examples/redblacktreeextended.go b/examples/redblacktreeextended.go index fbb8d6c..17855c9 100644 --- a/examples/redblacktreeextended.go +++ b/examples/redblacktreeextended.go @@ -95,6 +95,64 @@ func (tree *RedBlackTreeExtended) getMaxFromNode(node *rbt.Node) (foundNode *rbt } } +// Find ceiling node of the input key, return the ceiling node or nil if no ceiling 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 *RedBlackTreeExtended) Ceiling(key interface{}) (ceiling *rbt.Node, found bool) { + found = false + node := tree.Root + for node != nil { + compare := tree.Comparator(key, node.Key) + switch { + case compare == 0: + return node, true + case compare < 0: + ceiling, found = node, true + node = node.Left + case compare > 0: + node = node.Right + } + } + if found { + return ceiling, true + } + return nil, false +} + +// Find floor node of the input key, return the floor node or nil if no ceiling 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 *RedBlackTreeExtended) Floor(key interface{}) (floor *rbt.Node, found bool) { + found = false + node := tree.Root + for node != nil { + compare := tree.Comparator(key, node.Key) + switch { + case compare == 0: + return node, true + case compare < 0: + node = node.Left + case compare > 0: + floor, found = node, true + node = node.Right + } + } + if found { + return floor, true + } + return nil, false +} + func print(tree *RedBlackTreeExtended) { max, _ := tree.GetMax() min, _ := tree.GetMin() @@ -132,4 +190,30 @@ func RedBlackTreeExtendedExample() { // Value for min key: c // RedBlackTree // └── 3 + + // Ceiling and Floor functions + tree = RedBlackTreeExtended{rbt.NewWithIntComparator()} + tree.Put(1, "a") + tree.Put(2, "b") + tree.Put(4, "d") + tree.Put(6, "f") + tree.Put(7, "g") + + //index, ceiling, floor + testValues := [][]interface{}{ + {0, 1, nil}, + {1, 1, 1}, + {2, 2, 2}, + {3, 4, 2}, + {4, 4, 4}, + {5, 6, 4}, + {6, 6, 6}, + {7, 7, 7}, + {8, nil, 7}, + } + for _, tt := range testValues { + actualCeiling, _ := tree.Ceiling(tt[0]) + actualFloor, _ := tree.Floor(tt[0]) + fmt.Printf("test key %d, expected (%d, %d), actual (%d, %d)\n", tt[0], tt[1], tt[2], actualCeiling.Key, actualFloor.Key) + } } diff --git a/trees/redblacktree/redblacktree.go b/trees/redblacktree/redblacktree.go index 6e527fe..507b8d1 100644 --- a/trees/redblacktree/redblacktree.go +++ b/trees/redblacktree/redblacktree.go @@ -51,7 +51,7 @@ const ( type Tree struct { Root *Node size int - comparator utils.Comparator + Comparator utils.Comparator } type Node struct { @@ -65,17 +65,17 @@ type Node struct { // Instantiates a red-black tree with the custom comparator. func NewWith(comparator utils.Comparator) *Tree { - return &Tree{comparator: comparator} + return &Tree{Comparator: comparator} } // Instantiates a red-black tree with the IntComparator, i.e. keys are of type int. func NewWithIntComparator() *Tree { - return &Tree{comparator: utils.IntComparator} + return &Tree{Comparator: utils.IntComparator} } // Instantiates a red-black tree with the StringComparator, i.e. keys are of type string. func NewWithStringComparator() *Tree { - return &Tree{comparator: utils.StringComparator} + return &Tree{Comparator: utils.StringComparator} } // Inserts node into the tree. @@ -88,7 +88,7 @@ func (tree *Tree) Put(key interface{}, value interface{}) { node := tree.Root loop := true for loop { - compare := tree.comparator(key, node.Key) + compare := tree.Comparator(key, node.Key) switch { case compare == 0: node.Value = value @@ -263,7 +263,7 @@ func output(node *Node, prefix string, isTail bool, str *string) { func (tree *Tree) lookup(key interface{}) *Node { node := tree.Root for node != nil { - compare := tree.comparator(key, node.Key) + compare := tree.Comparator(key, node.Key) switch { case compare == 0: return node