From f7f48e7628930320e2807238c48776e495dcab77 Mon Sep 17 00:00:00 2001
From: otnt <jiangpufan@gmail.com>
Date: Sun, 27 Mar 2016 01:51:09 -0400
Subject: [PATCH] add ceiling and floor function to redblacktree

---
 trees/redblacktree/redblacktree.go      | 60 ++++++++++++++++++++++-
 trees/redblacktree/redblacktree_test.go | 64 +++++++++++++++++++++++++
 2 files changed, 123 insertions(+), 1 deletion(-)

diff --git a/trees/redblacktree/redblacktree.go b/trees/redblacktree/redblacktree.go
index 6e527fe..fd73c8a 100644
--- a/trees/redblacktree/redblacktree.go
+++ b/trees/redblacktree/redblacktree.go
@@ -116,7 +116,7 @@ func (tree *Tree) Put(key interface{}, value interface{}) {
 }
 
 // Searches the node in the tree by key and returns its value or nil if key is not found in tree.
-// Second return parameter is true if key was found, otherwise false.
+// return parameter is true if key was found, otherwise false.
 // Key should adhere to the comparator's type assertion, otherwise method panics.
 func (tree *Tree) Get(key interface{}) (value interface{}, found bool) {
 	node := tree.lookup(key)
@@ -126,6 +126,64 @@ 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{}) {
diff --git a/trees/redblacktree/redblacktree_test.go b/trees/redblacktree/redblacktree_test.go
index 58fe60f..af2f336 100644
--- a/trees/redblacktree/redblacktree_test.go
+++ b/trees/redblacktree/redblacktree_test.go
@@ -160,6 +160,70 @@ 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()