convert treemap in the same way as treeset

maps/treemap/enumerable.go

@@ -4,10 +4,7 @@
 
 package treemap
 
-import (
+import "github.com/emirpasic/gods/containers"
-	"github.com/emirpasic/gods/containers"
-	rbt "github.com/emirpasic/gods/trees/redblacktree"
-)
 
 func assertEnumerableImplementation() {
 	var _ containers.EnumerableWithKey = (*Map)(nil)
@@ -24,7 +21,7 @@ func (m *Map) Each(f func(key interface{}, value interface{})) {
 // Map invokes the given function once for each element and returns a container
 // containing the values returned by the given function as key/value pairs.
 func (m *Map) Map(f func(key1 interface{}, value1 interface{}) (interface{}, interface{})) *Map {
-	newMap := &Map{tree: rbt.NewWith(m.tree.Comparator)}
+	newMap := &Map{tree: m.tree.New()}
 	iterator := m.Iterator()
 	for iterator.Next() {
 		key2, value2 := f(iterator.Key(), iterator.Value())
@@ -35,7 +32,7 @@ func (m *Map) Map(f func(key1 interface{}, value1 interface{}) (interface{}, int
 
 // Select returns a new container containing all elements for which the given function returns a true value.
 func (m *Map) Select(f func(key interface{}, value interface{}) bool) *Map {
-	newMap := &Map{tree: rbt.NewWith(m.tree.Comparator)}
+	newMap := &Map{tree: m.tree.New()}
 	iterator := m.Iterator()
 	for iterator.Next() {
 		if f(iterator.Key(), iterator.Value()) {

maps/treemap/iterator.go

@@ -4,10 +4,7 @@
 
 package treemap
 
-import (
+import "github.com/emirpasic/gods/containers"
-	"github.com/emirpasic/gods/containers"
-	rbt "github.com/emirpasic/gods/trees/redblacktree"
-)
 
 func assertIteratorImplementation() {
 	var _ containers.ReverseIteratorWithKey = (*Iterator)(nil)
@@ -15,7 +12,7 @@ func assertIteratorImplementation() {
 
 // Iterator holding the iterator's state
 type Iterator struct {
-	iterator rbt.Iterator
+	iterator containers.ReverseIteratorWithKey
 }
 
 // Iterator returns a stateful iterator whose elements are key/value pairs.

maps/treemap/treemap.go

@@ -14,7 +14,8 @@ package treemap
 import (
 	"fmt"
 	"github.com/emirpasic/gods/maps"
-	rbt "github.com/emirpasic/gods/trees/redblacktree"
+	rbt "github.com/spewspews/gods/trees/redblacktree"
+	"github.com/spewspews/gods/trees"
 	"github.com/emirpasic/gods/utils"
 	"strings"
 )
@@ -25,7 +26,7 @@ func assertMapImplementation() {
 
 // Map holds the elements in a red-black tree
 type Map struct {
-	tree *rbt.Tree
+	tree trees.Tree
 }
 
 // NewWith instantiates a tree map with the custom comparator.
@@ -43,6 +44,11 @@ func NewWithStringComparator() *Map {
 	return &Map{tree: rbt.NewWithStringComparator()}
 }
 
+// NewWithTree instantiates a new empty map with given tree
+func NewWithTree(tree trees.Tree) *Map {
+	return &Map{tree: tree}
+}
+
 // Put inserts key-value pair into the map.
 // Key should adhere to the comparator's type assertion, otherwise method panics.
 func (m *Map) Put(key interface{}, value interface{}) {
@@ -90,19 +96,13 @@ func (m *Map) Clear() {
 // Min returns the minimum key and its value from the tree map.
 // Returns nil, nil if map is empty.
 func (m *Map) Min() (key interface{}, value interface{}) {
-	if node := m.tree.Left(); node != nil {
+	return m.tree.Min()
-		return node.Key, node.Value
-	}
-	return nil, nil
 }
 
 // Max returns the maximum key and its value from the tree map.
 // Returns nil, nil if map is empty.
 func (m *Map) Max() (key interface{}, value interface{}) {
-	if node := m.tree.Right(); node != nil {
+	return m.tree.Max()
-		return node.Key, node.Value
-	}
-	return nil, nil
 }
 
 // String returns a string representation of container

maps/treemap/treemap_test.go

@@ -7,6 +7,8 @@ package treemap
 import (
 	"fmt"
 	"testing"
+
+	"github.com/spewspews/gods/trees/avltree"
 )
 
 func TestMapPut(t *testing.T) {
@@ -51,6 +53,48 @@ func TestMapPut(t *testing.T) {
 	}
 }
 
+func TestMapAVLPut(t *testing.T) {
+	m := NewWithTree(avltree.NewWithIntComparator())
+	m.Put(5, "e")
+	m.Put(6, "f")
+	m.Put(7, "g")
+	m.Put(3, "c")
+	m.Put(4, "d")
+	m.Put(1, "x")
+	m.Put(2, "b")
+	m.Put(1, "a") //overwrite
+
+	if actualValue := m.Size(); actualValue != 7 {
+		t.Errorf("Got %v expected %v", actualValue, 7)
+	}
+	if actualValue, expectedValue := m.Keys(), []interface{}{1, 2, 3, 4, 5, 6, 7}; !sameElements(actualValue, expectedValue) {
+		t.Errorf("Got %v expected %v", actualValue, expectedValue)
+	}
+	if actualValue, expectedValue := m.Values(), []interface{}{"a", "b", "c", "d", "e", "f", "g"}; !sameElements(actualValue, expectedValue) {
+		t.Errorf("Got %v expected %v", actualValue, expectedValue)
+	}
+
+	// key,expectedValue,expectedFound
+	tests1 := [][]interface{}{
+		{1, "a", true},
+		{2, "b", true},
+		{3, "c", true},
+		{4, "d", true},
+		{5, "e", true},
+		{6, "f", true},
+		{7, "g", true},
+		{8, nil, false},
+	}
+
+	for _, test := range tests1 {
+		// retrievals
+		actualValue, actualFound := m.Get(test[0])
+		if actualValue != test[1] || actualFound != test[2] {
+			t.Errorf("Got %v expected %v", actualValue, test[1])
+		}
+	}
+
+}
 func TestMapRemove(t *testing.T) {
 	m := NewWithIntComparator()
 	m.Put(5, "e")

sets/treeset/treeset.go

@@ -46,7 +46,7 @@ func NewWithStringComparator() *Set {
 }
 
 // NewWithTree instantiates a new empty set with given tree
-func NewWithTree(tree trees.Tree) (set *Set) {
+func NewWithTree(tree trees.Tree) *Set {
 	return &Set{tree: tree}
 }
 

trees/avltree/avltree.go

@@ -386,6 +386,26 @@ func (t *Tree) Right() *Node {
 	return t.bottom(1)
 }
 
+// Min returns the minimum key value pair of the AVL tree
+// or nils if the tree is empty.
+func (t *Tree) Min() (interface{}, interface{}) {
+	n := t.bottom(0)
+	if n == nil {
+		return nil, nil
+	}
+	return n.Key, n.Value
+}
+
+// Max returns the minimum key value pair of the AVL tree
+// or nils if the tree is empty.
+func (t *Tree) Max() (interface{}, interface{}) {
+	n := t.bottom(1)
+	if n == nil {
+		return nil, nil
+	}
+	return n.Key, n.Value
+}
+
 func (t *Tree) bottom(d int) *Node {
 	n := t.Root
 	if n == nil {

trees/redblacktree/redblacktree.go

@@ -203,6 +203,26 @@ func (tree *Tree) Right() *Node {
 	return parent
 }
 
+// Min returns the minimum key value pair of the AVL tree
+// or nils if the tree is empty.
+func (t *Tree) Min() (interface{}, interface{}) {
+	n := t.Left()
+	if n == nil {
+		return nil, nil
+	}
+	return n.Key, n.Value
+}
+
+// Max returns the minimum key value pair of the AVL tree
+// or nils if the tree is empty.
+func (t *Tree) Max() (interface{}, interface{}) {
+	n := t.Right()
+	if n == nil {
+		return nil, nil
+	}
+	return n.Key, n.Value
+}
+
 // Floor Finds floor node of the input key, return the floor node or nil if no ceiling is found.
 // Second return parameter is true if floor was found, otherwise false.
 //

trees/trees.go

@@ -23,6 +23,8 @@ type Tree interface {
 	Remove(key interface{})
 	Get(key interface{}) (interface{}, bool)
 	Keys() []interface{}
+	Min() (interface{}, interface{})
+	Max() (interface{}, interface{})
 
 	containers.Container
 	// Empty() bool