完美avl树

avl/avl.go

@@ -47,7 +47,7 @@ func (avl *AVL) String() string {
 	if avl.size == 0 {
 		return ""
 	}
-	str := "AVL" + "\n"
+	str := "AVLTree" + "\n"
 	output(avl.root, "", true, &str)
 
 	return str

avl/avl_test.go

@@ -48,10 +48,34 @@ func TestIterator(t *testing.T) {
 	}
 	t.Error(avl.TraversalDepth(1))
 	iter := avl.Iterator()
+
+	for iter.Prev() {
+		t.Error(iter.Value())
+	}
+	t.Error("prev == false", iter.Value(), iter.Prev(), iter.Value())
+
 	for iter.Next() {
 		t.Error(iter.Value())
 	}
+	t.Error("next == false", iter.Value(), iter.Next(), iter.Value())
+
+	for iter.Prev() {
+		t.Error(iter.Value())
+	}
+	t.Error("prev == false", iter.Value())
+
+	for i := 0; iter.Next(); i++ {
+		t.Error(iter.Value())
+		if i >= 7 {
+			break
+		}
+	}
 	t.Error("next == false", iter.Value())
+
+	for iter.Prev() {
+		t.Error(iter.Value())
+	}
+	t.Error("prev == false", iter.Value())
 }
 
 func TestGetAround(t *testing.T) {
@@ -121,7 +145,9 @@ func TestPutComparatorRandom(t *testing.T) {
 		}
 
 		if avl.String() != godsavl.String() {
-			t.Error(content)
+			t.Error(godsavl.String())
+			t.Error(avl.String())
+			t.Error(content, n)
 			break
 		}
 	}
@@ -229,6 +255,58 @@ ALL:
 const PutCompartorSize = 300000
 const NumberMax = 60000000
 
+func BenchmarkIterator(b *testing.B) {
+	avl := New(utils.IntComparator)
+	b.N = PutCompartorSize
+
+	for i := 0; i < b.N; i++ {
+		v := randomdata.Number(0, NumberMax)
+		avl.Put(v)
+	}
+
+	b.ResetTimer()
+	b.StartTimer()
+	iter := avl.Iterator()
+	for iter.Next() {
+
+	}
+
+	for iter.Prev() {
+
+	}
+
+	for iter.Next() {
+
+	}
+
+}
+
+func BenchmarkGodsIterator(b *testing.B) {
+	avl := avltree.NewWithIntComparator()
+	b.N = PutCompartorSize
+
+	for i := 0; i < b.N; i++ {
+		v := randomdata.Number(0, NumberMax)
+		avl.Put(v, i)
+	}
+
+	b.ResetTimer()
+	b.StartTimer()
+	iter := avl.Iterator()
+	for iter.Next() {
+
+	}
+
+	for iter.Prev() {
+
+	}
+
+	for iter.Next() {
+
+	}
+
+}
+
 func BenchmarkRemove(b *testing.B) {
 
 	avl := New(utils.IntComparator)
@@ -316,6 +394,12 @@ func BenchmarkGodsAvlGet(b *testing.B) {
 func BenchmarkPut(b *testing.B) {
 	avl := New(utils.IntComparator)
 
+	for i := 0; i < 100000; i++ {
+		avl.Put(randomdata.Number(0, NumberMax))
+	}
+
+	b.ResetTimer()
+	b.StartTimer()
 	b.N = PutCompartorSize
 	for i := 0; i < b.N; i++ {
 		avl.Put(randomdata.Number(0, NumberMax))
@@ -326,6 +410,13 @@ func BenchmarkPut(b *testing.B) {
 func BenchmarkGodsRBPut(b *testing.B) {
 	rb := redblacktree.NewWithIntComparator()
 
+	for i := 0; i < 100000; i++ {
+		rb.Put(randomdata.Number(0, NumberMax), i)
+	}
+
+	b.ResetTimer()
+	b.StartTimer()
+
 	b.N = PutCompartorSize
 	for i := 0; i < b.N; i++ {
 		rb.Put(randomdata.Number(0, NumberMax), i)
@@ -336,6 +427,13 @@ func BenchmarkGodsRBPut(b *testing.B) {
 func BenchmarkGodsPut(b *testing.B) {
 	avl := avltree.NewWithIntComparator()
 
+	for i := 0; i < 100000; i++ {
+		avl.Put(randomdata.Number(0, NumberMax), i)
+	}
+
+	b.ResetTimer()
+	b.StartTimer()
+
 	b.N = PutCompartorSize
 	for i := 0; i < b.N; i++ {
 		avl.Put(randomdata.Number(0, NumberMax), i)

avl/iterator.go

@@ -1,24 +1,22 @@
 package avl
 
 import (
-	"log"
-
 	"github.com/emirpasic/gods/stacks/arraystack"
 )
 
 type Iterator struct {
 	op *AVL
 
-	cur       *Node
+	dir    int
-	prevstack *arraystack.Stack
+	up     *Node
-	nextstack *arraystack.Stack
+	cur    *Node
+	tstack *arraystack.Stack
 	// curnext *Node
 }
 
 func initIterator(avltree *AVL) *Iterator {
-	iter := &Iterator{op: avltree, prevstack: arraystack.New(), nextstack: arraystack.New()}
+	iter := &Iterator{op: avltree, tstack: arraystack.New()}
-	iter.cur = avltree.root
+	iter.up = avltree.root
-	// iter.nextstack.Push(avltree.root)
 	return iter
 }
 
@@ -28,14 +26,46 @@ func (iter *Iterator) Value() interface{} {
 
 func (iter *Iterator) Next() (result bool) {
 
-	if iter.nextstack.Size() == 0 {
+	if iter.dir < 1 {
-		iter.nextstack.Push(iter.cur)
+		if iter.dir == -1 {
+			if iter.cur != nil {
+				iter.tstack.Clear()
+				iter.curPushNextStack(iter.cur)
+
+				iter.up = iter.cur
+				for iter.up != nil {
+					if iter.up.child == 1 {
+						iter.up = iter.up.parent
+						break
+					} else {
+						iter.up = iter.up.parent
+					}
+				}
+
+			}
+		}
+		iter.dir = 1
 	}
 
-	if v, ok := iter.nextstack.Pop(); ok {
+	if iter.tstack.Size() == 0 {
+		if iter.up != nil {
+			iter.tstack.Push(iter.up)
+			for iter.up != nil {
+				if iter.up.child == 1 {
+					iter.up = iter.up.parent
+					break
+				} else {
+					iter.up = iter.up.parent
+				}
+			}
+		} else {
+			return false
+		}
+	}
+
+	if v, ok := iter.tstack.Pop(); ok {
 		iter.cur = v.(*Node)
 		iter.curPushNextStack(iter.cur)
-		log.Println(iter.nextstack.Size())
 		return true
 	}
 
@@ -46,19 +76,68 @@ func (iter *Iterator) curPushNextStack(cur *Node) {
 	next := cur.children[0]
 
 	if next != nil {
-		iter.nextstack.Push(next)
+		iter.tstack.Push(next)
 		for next.children[1] != nil {
 			next = next.children[1]
-			iter.nextstack.Push(next)
+			iter.tstack.Push(next)
+		}
+	}
+}
+
+func (iter *Iterator) curPushPrevStack(cur *Node) {
+	prev := cur.children[1]
+
+	if prev != nil {
+		iter.tstack.Push(prev)
+		for prev.children[0] != nil {
+			prev = prev.children[0]
+			iter.tstack.Push(prev)
 		}
 	}
 }
 
 func (iter *Iterator) Prev() (result bool) {
 
-	if v, ok := iter.nextstack.Pop(); ok {
+	if iter.dir > -1 {
+
+		if iter.dir == 1 {
+			if iter.cur != nil {
+				iter.tstack.Clear()
+				iter.curPushPrevStack(iter.cur)
+
+				iter.up = iter.cur
+				for iter.up != nil {
+					if iter.up.child == 0 {
+						iter.up = iter.up.parent
+						break
+					} else {
+						iter.up = iter.up.parent
+					}
+				}
+			}
+		}
+		iter.dir = -1
+	}
+
+	if iter.tstack.Size() == 0 {
+		if iter.up != nil {
+			iter.tstack.Push(iter.up)
+			for iter.up != nil {
+				if iter.up.child == 0 {
+					iter.up = iter.up.parent
+					break
+				} else {
+					iter.up = iter.up.parent
+				}
+			}
+		} else {
+			return false
+		}
+	}
+
+	if v, ok := iter.tstack.Pop(); ok {
 		iter.cur = v.(*Node)
-		iter.curPushNextStack(iter.cur)
+		iter.curPushPrevStack(iter.cur)
 		return true
 	}