diff --git a/examples/redblacktreeextended.go b/examples/redblacktreeextended.go index fbb8d6c..df148f3 100644 --- a/examples/redblacktreeextended.go +++ b/examples/redblacktreeextended.go @@ -103,6 +103,66 @@ func print(tree *RedBlackTreeExtended) { fmt.Println(tree) } +// Find ceiling node of the input key, return its key and value or nil if no ceiling is found. +// Third 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{}) (ceilingKey interface{}, value interface{}, found bool) { + var ceiling *rbt.Node + found = false + comparator := tree.Comparator() + + node := tree.Root + for node != nil { + compare := comparator(key, node.Key) + switch { + case compare == 0: + return node.Key, node.Value, true + case compare < 0: + ceiling, found = node, true + node = node.Left + case compare > 0: + node = node.Right + } + } + if found { + return ceiling.Key, ceiling.Value, true + } + return nil, nil, false +} + +// Find floor node of the input key, return its key and value or nil if no ceiling is found. +// Third 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{}) (floorKey interface{}, value interface{}, found bool) { + var floor *rbt.Node + found = false + comparator := tree.Comparator() + + node := tree.Root + for node != nil { + compare := comparator(key, node.Key) + switch { + case compare == 0: + return node.Key, node.Value, true + case compare < 0: + node = node.Left + case compare > 0: + floor, found = node, true + node = node.Right + } + } + if found { + return floor.Key, floor.Value, true + } + return nil, nil, false +} + func RedBlackTreeExtendedExample() { tree := RedBlackTreeExtended{rbt.NewWithIntComparator()} @@ -132,4 +192,32 @@ 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, actualFloor) + } } diff --git a/trees/redblacktree/redblacktree.go b/trees/redblacktree/redblacktree.go index 4ef4d14..fe56a6c 100644 --- a/trees/redblacktree/redblacktree.go +++ b/trees/redblacktree/redblacktree.go @@ -126,64 +126,6 @@ func (tree *Tree) Get(key interface{}) (value interface{}, found bool) { return nil, false } -// Find ceiling node of the input key, return its key and value or nil if no ceiling is found. -// Third 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 *Tree) Ceiling(key interface{}) (ceilingKey interface{}, value interface{}, found bool) { - ceiling := &Node{} - found = false - - node := tree.Root - for node != nil { - compare := tree.comparator(key, node.Key) - switch { - case compare == 0: - return node.Key, node.Value, true - case compare < 0: - ceiling, found = node, true - node = node.Left - case compare > 0: - node = node.Right - } - } - if found { - return ceiling.Key, ceiling.Value, true - } - return nil, nil, false -} - -// Find floor node of the input key, return its key and value or nil if no ceiling is found. -// Third 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 *Tree) Floor(key interface{}) (floorKey interface{}, value interface{}, found bool) { - floor := &Node{} - found = false - - node := tree.Root - for node != nil { - compare := tree.comparator(key, node.Key) - switch { - case compare == 0: - return node.Key, node.Value, true - case compare < 0: - node = node.Left - case compare > 0: - floor, found = node, true - node = node.Right - } - } - if found { - return floor.Key, floor.Value, true - } - return nil, nil, false -} - // Remove the node from the tree by key. // Key should adhere to the comparator's type assertion, otherwise method panics. func (tree *Tree) Remove(key interface{}) { @@ -250,6 +192,11 @@ func (tree *Tree) Clear() { tree.size = 0 } +// Return comparator of the tree +func (tree *Tree) Comparator() utils.Comparator { + return tree.comparator +} + func (tree *Tree) String() string { str := "RedBlackTree\n" if !tree.Empty() { diff --git a/trees/redblacktree/redblacktree_test.go b/trees/redblacktree/redblacktree_test.go index af2f336..58fe60f 100644 --- a/trees/redblacktree/redblacktree_test.go +++ b/trees/redblacktree/redblacktree_test.go @@ -160,70 +160,6 @@ func TestRedBlackTree(t *testing.T) { } -func TestCeiling(t *testing.T) { - - tree := NewWithIntComparator() - - //key-value to insert - tree.Put(1, "a") - tree.Put(2, "b") - tree.Put(4, "d") - tree.Put(6, "f") - tree.Put(7, "g") - - //ceiling map - ceilingMap := [][]interface{}{ - {0, 1, true}, - {1, 1, true}, - {2, 2, true}, - {3, 4, true}, - {4, 4, true}, - {5, 6, true}, - {6, 6, true}, - {7, 7, true}, - {8, nil, false}, - } - for _, test := range ceilingMap { - actualKey, _, actualFound := tree.Ceiling(test[0]) - if actualKey != test[1] || actualFound != test[2] { - t.Errorf("Got (%v, %v) expected (%v, %v)", - actualKey, actualFound, test[1], test[2]) - } - } -} - -func TestFloor(t *testing.T) { - - tree := NewWithIntComparator() - - //key-value to insert - tree.Put(1, "a") - tree.Put(2, "b") - tree.Put(4, "d") - tree.Put(6, "f") - tree.Put(7, "g") - - //ceiling map - ceilingMap := [][]interface{}{ - {0, nil, false}, - {1, 1, true}, - {2, 2, true}, - {3, 2, true}, - {4, 4, true}, - {5, 4, true}, - {6, 6, true}, - {7, 7, true}, - {8, 7, true}, - } - for _, test := range ceilingMap { - actualKey, _, actualFound := tree.Floor(test[0]) - if actualKey != test[1] || actualFound != test[2] { - t.Errorf("Got (%v, %v) expected (%v, %v)", - actualKey, actualFound, test[1], test[2]) - } - } -} - func BenchmarkRedBlackTree(b *testing.B) { for i := 0; i < b.N; i++ { tree := NewWithIntComparator()