From 1078d101894d0948bb488cb34206282d715aab68 Mon Sep 17 00:00:00 2001
From: loxp <loxp1991@gmail.com>
Date: Wed, 15 Aug 2018 17:23:51 +0800
Subject: [PATCH] Add Floor and Ceiling method to Treemap

---
 maps/treemap/treemap.go      | 22 +++++++++++++++
 maps/treemap/treemap_test.go | 54 ++++++++++++++++++++++++++++++++++++
 2 files changed, 76 insertions(+)

diff --git a/maps/treemap/treemap.go b/maps/treemap/treemap.go
index a1e58ad..2105b66 100644
--- a/maps/treemap/treemap.go
+++ b/maps/treemap/treemap.go
@@ -105,6 +105,28 @@ func (m *Map) Max() (key interface{}, value interface{}) {
 	return nil, nil
 }
 
+// Floor searches the floor element in the map by key.
+// Returns floor key, floor value, true if floor key is found.
+// Returns nil, nil, false if floor key is not found.
+func (m *Map) Floor(key interface{}) (retKey interface{}, retValue interface{}, found bool) {
+	ret, found := m.tree.Floor(key)
+	if !found {
+		return nil, nil, false
+	}
+	return ret.Key, ret.Value, true
+}
+
+// Ceiling searches the ceiling element in the map by key.
+// Returns ceiling key, ceiling value, true if a ceiling key is found.
+// Returns nil, nil, false if ceiling key is not found.
+func (m *Map) Ceiling(key interface{}) (retKey interface{}, retValue interface{}, found bool) {
+	ret, found := m.tree.Ceiling(key)
+	if !found {
+		return nil, nil, false
+	}
+	return ret.Key, ret.Value, true
+}
+
 // String returns a string representation of container
 func (m *Map) String() string {
 	str := "TreeMap\nmap["
diff --git a/maps/treemap/treemap_test.go b/maps/treemap/treemap_test.go
index a73e873..da82bdc 100644
--- a/maps/treemap/treemap_test.go
+++ b/maps/treemap/treemap_test.go
@@ -289,6 +289,60 @@ func TestMapChaining(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, actualFound := m.Floor(test[0])
+		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, actualFound := m.Ceiling(test[0])
+		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 TestMapIteratorNextOnEmpty(t *testing.T) {
 	m := NewWithStringComparator()
 	it := m.Iterator()