TODO: GetRange

avl/avl.go

@@ -120,6 +120,19 @@ func (avl *Tree) Remove(key interface{}) *Node {
 	return nil
 }
 
+// Values 返回先序遍历的值
+func (avl *Tree) Values() []interface{} {
+	mszie := 0
+	if avl.root != nil {
+		mszie = avl.size
+	}
+	result := make([]interface{}, 0, mszie)
+	avl.Traversal(func(v interface{}) {
+		result = append(result, v)
+	}, DLR)
+	return result
+}
+
 func (avl *Tree) Get(key interface{}) (interface{}, bool) {
 	n, ok := avl.GetNode(key)
 	if ok {
@@ -143,57 +156,35 @@ func (avl *Tree) GetAround(key interface{}) (result [3]interface{}) {
 }
 
 func (avl *Tree) GetAroundNode(value interface{}) (result [3]*Node) {
-	n := avl.root
+	if cur, ok := avl.GetNode(value); ok {
 
-	for {
+		var iter *Iterator
 
-		if n == nil {
-			return
+		iter = NewIterator(cur)
+		iter.curPushPrevStack(iter.up)
+		iter.up = iter.getPrevUp(iter.up)
+
+		if v, ok := iter.tstack.Pop(); ok {
+			result[0] = v.(*Node)
+			// iter.curPushPrevStack(iter.cur)
+		} else {
+			result[0] = iter.up
 		}
 
-		lastc := 0
-		switch c := avl.comparator(value, n.value); c {
-		case -1:
-			if c != -lastc {
-				result[0] = n
-			}
-			lastc = c
-			n = n.children[0]
-		case 1:
-			if c != -lastc {
-				result[2] = n
-			}
-			lastc = c
-			n = n.children[1]
-		case 0:
+		iter = NewIterator(cur)
+		iter.curPushNextStack(iter.up)
+		iter.up = iter.getNextUp(iter.up)
 
-			switch lastc {
-			case -1:
-				if n.children[1] != nil {
-					result[0] = n.children[1]
-				}
-			case 1:
-				if n.children[0] != nil {
-					result[2] = n.children[0]
-				}
-			case 0:
-
-				if n.children[1] != nil {
-					result[0] = n.children[1]
-				}
-				if n.children[0] != nil {
-					result[2] = n.children[0]
-				}
-
-				result[1] = n
-				return
-			}
-
-		default:
-			panic("Get comparator only is allowed in -1, 0, 1")
+		if v, ok := iter.tstack.Pop(); ok {
+			result[2] = v.(*Node)
+			// iter.curPushNextStack(iter.cur)
+		} else {
+			result[2] = iter.up
 		}
 
+		result[1] = cur
 	}
+	return
 }
 func (avl *Tree) GetNode(value interface{}) (*Node, bool) {
 
@@ -251,47 +242,64 @@ func (avl *Tree) debugString() string {
 	return str
 }
 
-func (avl *Tree) TraversalBreadth() (result []interface{}) {
-	result = make([]interface{}, 0, avl.size)
-	var traverasl func(cur *Node)
-	traverasl = func(cur *Node) {
-		if cur == nil {
-			return
-		}
-		result = append(result, cur.value)
-		traverasl(cur.children[0])
-		traverasl(cur.children[1])
-	}
-	traverasl(avl.root)
-	return
-}
+type TraversalMethod int
 
-func (avl *Tree) TraversalDepth(leftright int) (result []interface{}) {
-	result = make([]interface{}, 0, avl.size)
-	if leftright < 0 {
+const (
+	_ TraversalMethod = iota
+	//DLR 前序遍历
+	DLR
+	//LDR 中序遍历
+	LDR
+	//LRD 后序遍历
+	LRD
+)
+
+// Traversal 遍历的方法
+func (avl *Tree) Traversal(every func(v interface{}), traversalMethod ...interface{}) {
+	if avl.root == nil {
+		return
+	}
+
+	method := LDR
+	if len(traversalMethod) != 0 {
+		method = traversalMethod[0].(TraversalMethod)
+	}
+
+	switch method {
+	case DLR:
 		var traverasl func(cur *Node)
 		traverasl = func(cur *Node) {
 			if cur == nil {
 				return
 			}
 			traverasl(cur.children[0])
-			result = append(result, cur.value)
+			every(cur.value)
 			traverasl(cur.children[1])
 		}
 		traverasl(avl.root)
-	} else {
+	case LRD:
 		var traverasl func(cur *Node)
 		traverasl = func(cur *Node) {
 			if cur == nil {
 				return
 			}
 			traverasl(cur.children[1])
-			result = append(result, cur.value)
+			every(cur.value)
 			traverasl(cur.children[0])
 		}
 		traverasl(avl.root)
+	case LDR:
+		var traverasl func(cur *Node)
+		traverasl = func(cur *Node) {
+			if cur == nil {
+				return
+			}
+			every(cur.value)
+			traverasl(cur.children[0])
+			traverasl(cur.children[1])
+		}
+		traverasl(avl.root)
 	}
-
 	return
 }
 
@@ -564,7 +572,6 @@ func abs(n int) int {
 }
 
 func (avl *Tree) fixRemoveHeight(cur *Node) {
-
 	for {
 
 		lefth, rigthh, lrmax := getMaxAndChildrenHeight(cur)

avl/avl_test.go

@@ -16,7 +16,7 @@ import (
 	"github.com/emirpasic/gods/utils"
 )
 
-const CompartorSize = 1000000
+const CompartorSize = 100
 const NumberMax = 600
 
 func TestSave(t *testing.T) {
@@ -124,7 +124,7 @@ func TestGetAround(t *testing.T) {
 	}
 
 	if spew.Sprint(avl.GetAround(40)) != "[40 40 30]" {
-		t.Error("avl.GetAround(40)) is error", spew.Sprint(avl.GetAround(50)))
+		t.Error("avl.GetAround(40)) is error", spew.Sprint(avl.GetAround(40)))
 	}
 
 	if spew.Sprint(avl.GetAround(50)) != "[<nil> 50 40]" {
@@ -241,7 +241,7 @@ ALL:
 		for i := 0; i < 100; i++ {
 			avl.Remove(l[i])
 			gods.Remove(l[i])
-			s1 := spew.Sprint(avl.TraversalDepth(-1))
+			s1 := spew.Sprint(avl.Values())
 			s2 := spew.Sprint(gods.Values())
 			if s1 != s2 {
 				t.Error("avl remove error", "avlsize = ", avl.Size())
@@ -251,6 +251,7 @@ ALL:
 			}
 		}
 	}
+
 }
 
 func TestRemove(t *testing.T) {
@@ -279,7 +280,7 @@ ALL:
 		for i := 0; i < 10; i++ {
 			avl.Remove(l[i])
 			gods.Remove(l[i])
-			if spew.Sprint(gods.Values()) != spew.Sprint(avl.TraversalDepth(-1)) && avl.size != 0 {
+			if spew.Sprint(gods.Values()) != spew.Sprint(avl.Values()) && avl.size != 0 {
 				// if gods.String() != avl.String() && gods.Size() != 0 && avl.size != 0 {
 				t.Error(src1)
 				t.Error(src2)
@@ -316,27 +317,6 @@ func BenchmarkIterator(b *testing.B) {
 
 }
 
-func BenchmarkGodsIterator(b *testing.B) {
-	tree := avltree.NewWithIntComparator()
-
-	l := loadTestData()
-	b.N = len(l)
-
-	for _, v := range l {
-		tree.Put(v, v)
-	}
-
-	b.ResetTimer()
-	b.StartTimer()
-	iter := tree.Iterator()
-	for iter.Next() {
-	}
-	for iter.Prev() {
-	}
-	for iter.Next() {
-	}
-}
-
 func BenchmarkRemove(b *testing.B) {
 	tree := New(utils.IntComparator)
 

avl/iterator.go

@@ -20,8 +20,10 @@ func initIterator(avltree *Tree) *Iterator {
 	return iter
 }
 
-func NewIterator(tree *Tree) *Iterator {
-	return initIterator(tree)
+func NewIterator(n *Node) *Iterator {
+	iter := &Iterator{tstack: lastack.New()}
+	iter.up = n
+	return iter
 }
 
 func (iter *Iterator) Value() interface{} {