move Ceiling and Floor function to redblacktreeextend

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

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() {

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()