TODO: iter call 太多系统内存. 导致效率低下

avl/avl.go

@@ -28,10 +28,11 @@ type Tree struct {
 	root    *Node
 	size    int
 	Compare compare.Compare
+	iter    *Iterator
 }
 
 func New(Compare compare.Compare) *Tree {
-	return &Tree{Compare: Compare}
+	return &Tree{Compare: Compare, iter: NewIteratorWithCap(nil, 16)}
 }
 
 func (tree *Tree) String() string {
@@ -285,11 +286,12 @@ func (tree *Tree) GetNode(value interface{}) (*Node, bool) {
 		case 1:
 			n = n.children[1]
 		case 0:
-			iter := NewIterator(n)
-			iter.Prev()
-			for iter.Prev() {
-				if tree.Compare(iter.cur.value, n.value) == 0 {
-					n = iter.cur
+			// iter := NewIterator(n)
+			tree.iter.SetNode(n)
+			tree.iter.Prev()
+			for tree.iter.Prev() {
+				if tree.Compare(tree.iter.cur.value, n.value) == 0 {
+					n = tree.iter.cur
 				} else {
 					break
 				}
@@ -310,25 +312,29 @@ func (tree *Tree) Put(value interface{}) {
 		return
 	}
 
-	cur := tree.root
-	parent := cur.parent
-	child := -1
-
-	for {
-
-		if cur == nil {
-			parent.children[child] = node
-			node.parent = parent
-			if node.parent.height == 0 {
-				tree.fixPutHeight(node.parent)
+	for cur, c := tree.root, 0; ; {
+		c = tree.Compare(value, cur.value)
+		if c == -1 {
+			if cur.children[0] == nil {
+				cur.children[0] = node
+				node.parent = cur
+				if cur.height == 0 {
+					tree.fixPutHeight(cur)
+				}
+				return
 			}
-			return
+			cur = cur.children[0]
+		} else {
+			if cur.children[1] == nil {
+				cur.children[1] = node
+				node.parent = cur
+				if cur.height == 0 {
+					tree.fixPutHeight(cur)
+				}
+				return
+			}
+			cur = cur.children[1]
 		}
-
-		parent = cur
-		c := tree.Compare(value, cur.value)
-		child = (c + 2) / 2
-		cur = cur.children[child]
 	}
 }
 

avl/avl_test.go

@@ -247,11 +247,11 @@ func TestPutStable(t *testing.T) {
 	// log.SetOutput(f)
 	// 0-1 3 | 2-3 7-8 | 4-7 12-16 | 8-15 20-32 | 16-31 33-58 l := []int{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 18, 19, 20, 21, 22, 30, 41, 41, 41}
 
-	tree := New(compare.Int)
-	for i := 0; i < 10; i++ {
-		tree.Put(randomdata.Number(0, 100))
-	}
-	t.Error(tree.debugString())
+	// tree := New(compare.Int)
+	// for i := 0; i < 10; i++ {
+	// 	tree.Put(randomdata.Number(0, 100))
+	// }
+	// t.Error(tree.debugString())
 
 	// t.Error(tree.debugString(), tree.TraversalBreadth(), "\n", "-----------")
 }
@@ -485,11 +485,20 @@ func BenchmarkGet(b *testing.B) {
 
 	l := loadTestData()
 	b.N = len(l)
+	for i := 0; i < b.N; i++ {
+		tree.Put(l[i])
+	}
 
 	b.ResetTimer()
 	b.StartTimer()
-	for i := 0; i < b.N; i++ {
-		tree.Get(l[i])
+
+	execCount := 50
+	b.N = len(l) * execCount
+
+	for i := 0; i < execCount; i++ {
+		for _, v := range l {
+			tree.Get(v)
+		}
 	}
 }
 
@@ -526,7 +535,7 @@ func BenchmarkPut(b *testing.B) {
 	b.ResetTimer()
 	b.StartTimer()
 
-	execCount := 50
+	execCount := 1000
 	b.N = len(l) * execCount
 	for i := 0; i < execCount; i++ {
 		tree := New(compare.Int)

avl/iterator.go

@@ -24,6 +24,17 @@ func NewIterator(n *Node) *Iterator {
 	return iter
 }
 
+func NewIteratorWithCap(n *Node, cap int) *Iterator {
+	iter := &Iterator{tstack: lastack.NewWithCap(cap)}
+	iter.up = n
+	return iter
+}
+
+func (iter *Iterator) SetNode(n *Node) {
+	iter.up = n
+	iter.tstack.Clear()
+}
+
 func (iter *Iterator) Value() interface{} {
 	return iter.cur.value
 }