TreeMap Floor and Ceiling functions

maps/treemap/treemap.go

@@ -105,6 +105,40 @@ func (m *Map) Max() (key interface{}, value interface{}) {
 	return nil, nil
 }
 
+// Floor finds the floor key-value pair for the input key.
+// In case that no floor is found, then both returned values will be nil.
+// It's generally enough to check the first value (key) for nil, which determines if floor was found.
+//
+// Floor key is defined as the largest key that is smaller than or equal to the given key.
+// A floor key may not be found, either because the map is empty, or because
+// all keys in the map are larger than the given key.
+//
+// Key should adhere to the comparator's type assertion, otherwise method panics.
+func (m *Map) Floor(key interface{}) (foundKey interface{}, foundValue interface{}) {
+	node, found := m.tree.Floor(key)
+	if found {
+		return node.Key, node.Value
+	}
+	return nil, nil
+}
+
+// Ceiling finds the ceiling key-value pair for the input key.
+// In case that no ceiling is found, then both returned values will be nil.
+// It's generally enough to check the first value (key) for nil, which determines if ceiling was found.
+//
+// Ceiling key is defined as the smallest key that is larger than or equal to the given key.
+// A ceiling key may not be found, either because the map is empty, or because
+// all keys in the map are smaller than the given key.
+//
+// Key should adhere to the comparator's type assertion, otherwise method panics.
+func (m *Map) Ceiling(key interface{}) (foundKey interface{}, foundValue interface{}) {
+	node, found := m.tree.Ceiling(key)
+	if found {
+		return node.Key, node.Value
+	}
+	return nil, nil
+}
+
 // String returns a string representation of container
 func (m *Map) String() string {
 	str := "TreeMap\nmap["

maps/treemap/treemap_test.go

@@ -118,6 +118,62 @@ func TestMapRemove(t *testing.T) {
 	}
 }
 
+func TestMapFloor(t *testing.T) {
+	m := NewWithIntComparator()
+	m.Put(7, "g")
+	m.Put(3, "c")
+	m.Put(1, "a")
+
+	// key,expectedKey,expectedValue,expectedFound
+	tests1 := [][]interface{}{
+		{-1, nil, nil, false},
+		{0, nil, nil, false},
+		{1, 1, "a", true},
+		{2, 1, "a", true},
+		{3, 3, "c", true},
+		{4, 3, "c", true},
+		{7, 7, "g", true},
+		{8, 7, "g", true},
+	}
+
+	for _, test := range tests1 {
+		// retrievals
+		actualKey, actualValue := m.Floor(test[0])
+		actualFound := actualKey != nil && actualValue != nil
+		if actualKey != test[1] || actualValue != test[2] || actualFound != test[3] {
+			t.Errorf("Got %v, %v, %v, expected %v, %v, %v", actualKey, actualValue, actualFound, test[1], test[2], test[3])
+		}
+	}
+}
+
+func TestMapCeiling(t *testing.T) {
+	m := NewWithIntComparator()
+	m.Put(7, "g")
+	m.Put(3, "c")
+	m.Put(1, "a")
+
+	// key,expectedKey,expectedValue,expectedFound
+	tests1 := [][]interface{}{
+		{-1, 1, "a", true},
+		{0, 1, "a", true},
+		{1, 1, "a", true},
+		{2, 3, "c", true},
+		{3, 3, "c", true},
+		{4, 7, "g", true},
+		{7, 7, "g", true},
+		{8, nil, nil, false},
+	}
+
+	for _, test := range tests1 {
+		// retrievals
+		actualKey, actualValue := m.Ceiling(test[0])
+		actualFound := actualKey != nil && actualValue != nil
+		if actualKey != test[1] || actualValue != test[2] || actualFound != test[3] {
+			t.Errorf("Got %v, %v, %v, expected %v, %v, %v", actualKey, actualValue, actualFound, test[1], test[2], test[3])
+		}
+	}
+}
+
 func sameElements(a []interface{}, b []interface{}) bool {
 	if len(a) != len(b) {
 		return false

trees/redblacktree/redblacktree.go

@@ -197,12 +197,12 @@ func (tree *Tree) Right() *Node {
 	return parent
 }
 
-// Floor Finds floor node of the input key, return the floor node or nil if no ceiling is found.
+// Floor Finds floor node of the input key, return the floor node or nil if no floor 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.
+// all nodes in the tree are larger than the given node.
 //
 // Key should adhere to the comparator's type assertion, otherwise method panics.
 func (tree *Tree) Floor(key interface{}) (floor *Node, found bool) {
@@ -231,7 +231,7 @@ func (tree *Tree) Floor(key interface{}) (floor *Node, found bool) {
 //
 // 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.
+// all nodes in the tree are smaller than the given node.
 //
 // Key should adhere to the comparator's type assertion, otherwise method panics.
 func (tree *Tree) Ceiling(key interface{}) (ceiling *Node, found bool) {