From 09dbba1a881ac99bc1bd9cdd8558a845d394e039 Mon Sep 17 00:00:00 2001
From: eson <474420502@qq.com>
Date: Sun, 7 Apr 2019 05:31:48 +0800
Subject: [PATCH] =?UTF-8?q?TODO:=20fixSizeWithRemove=20=E5=8F=AF=E8=83=BD?=
 =?UTF-8?q?=E4=B8=BAnil?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 priority_queue/vbt.go      | 278 +++++++++++++++----------------------
 priority_queue/vbt_test.go |  14 +-
 2 files changed, 116 insertions(+), 176 deletions(-)

diff --git a/priority_queue/vbt.go b/priority_queue/vbt.go
index 4873659..272a0c1 100644
--- a/priority_queue/vbt.go
+++ b/priority_queue/vbt.go
@@ -439,7 +439,7 @@ func (tree *vbTree) Put(key interface{}) {
 		if cur.size > 8 {
 			factor := cur.size / 10 // or factor = 1
 			if cur.children[1].size >= cur.children[0].size*2+factor || cur.children[0].size >= cur.children[1].size*2+factor {
-				tree.fixSize(cur)
+				cur = tree.fixSize(cur)
 			}
 		}
 
@@ -624,22 +624,29 @@ func (tree *vbTree) Traversal(every func(v interface{}) bool, traversalMethod ..
 	}
 }
 
-func setChild(cur *tNode, cidx int, child *tNode) {
+func setChildNotNil(cur *tNode, cidx int, child *tNode) {
 	cur.children[cidx] = child
 	cur.children[cidx].parent = cur
 }
 
-func (tree *vbTree) replaceParent(old, new *tNode) {
-	if old.parent == nil {
-		tree.root = new
+func setChild(cur *tNode, cidx int, child *tNode) {
+	cur.children[cidx] = child
+	if child != nil {
+		cur.children[cidx].parent = cur
+	}
+}
+
+func (tree *vbTree) takeParent(token, person *tNode) {
+	if token.parent == nil {
+		tree.root = person
 	} else {
-		if old.parent.children[1] == old {
-			old.parent.children[1] = new
+		if token.parent.children[1] == token {
+			token.parent.children[1] = person
 		} else {
-			old.parent.children[0] = new
+			token.parent.children[0] = person
 		}
 	}
-	new.parent = old.parent
+	person.parent = token.parent
 }
 
 func (tree *vbTree) lrrotate3(cur *tNode) *tNode {
@@ -652,9 +659,9 @@ func (tree *vbTree) lrrotate3(cur *tNode) *tNode {
 	lrn := ln.children[r]
 	ln.children[r] = nil
 
-	tree.replaceParent(cur, lrn)
-	setChild(lrn, l, ln)
-	setChild(lrn, r, cur)
+	tree.takeParent(cur, lrn)
+	setChildNotNil(lrn, l, ln)
+	setChildNotNil(lrn, r, cur)
 
 	lrn.size = 3
 	lrn.children[l].size = 1
@@ -662,44 +669,30 @@ func (tree *vbTree) lrrotate3(cur *tNode) *tNode {
 	return lrn
 }
 
-func (tree *vbTree) lrrotate(cur *tNode) {
+func (tree *vbTree) lrrotate(cur *tNode) *tNode {
 
 	const l = 1
 	const r = 0
 
-	movparent := cur.children[l]
-	mov := movparent.children[r]
+	ln := cur.children[l]
+	lrn := ln.children[r]
 
-	mov.value, cur.value = cur.value, mov.value //交换值达到, 相对位移
+	lrln := lrn.children[l]
+	lrrn := lrn.children[r]
 
-	if mov.children[l] != nil {
-		movparent.children[r] = mov.children[l]
-		movparent.children[r].parent = movparent
-		//movparent.children[r].child = l
-	} else {
-		movparent.children[r] = nil
-	}
+	tree.takeParent(cur, lrn)
 
-	if mov.children[r] != nil {
-		mov.children[l] = mov.children[r]
-		//mov.children[l].child = l
-	} else {
-		mov.children[l] = nil
-	}
+	setChild(ln, r, lrln)
+	setChild(cur, l, lrrn)
 
-	if cur.children[r] != nil {
-		mov.children[r] = cur.children[r]
-		mov.children[r].parent = mov
-	} else {
-		mov.children[r] = nil
-	}
+	setChildNotNil(lrn, l, ln)
+	setChildNotNil(lrn, r, cur)
 
-	cur.children[r] = mov
-	mov.parent = cur
-
-	movparent.size = getChildrenSumSize(movparent) + 1
-	mov.size = getChildrenSumSize(mov) + 1
+	ln.size = getChildrenSumSize(ln) + 1
 	cur.size = getChildrenSumSize(cur) + 1
+	lrn.size = getChildrenSumSize(lrn) + 1
+
+	return lrn
 }
 
 func (tree *vbTree) rlrotate3(cur *tNode) *tNode {
@@ -712,9 +705,9 @@ func (tree *vbTree) rlrotate3(cur *tNode) *tNode {
 	lrn := ln.children[r]
 	ln.children[r] = nil
 
-	tree.replaceParent(cur, lrn)
-	setChild(lrn, l, ln)
-	setChild(lrn, r, cur)
+	tree.takeParent(cur, lrn)
+	setChildNotNil(lrn, l, ln)
+	setChildNotNil(lrn, r, cur)
 
 	lrn.size = 3
 	lrn.children[l].size = 1
@@ -722,54 +715,62 @@ func (tree *vbTree) rlrotate3(cur *tNode) *tNode {
 	return lrn
 }
 
-func (tree *vbTree) rlrotate(cur *tNode) {
+func (tree *vbTree) rlrotate(cur *tNode) *tNode {
 
 	const l = 0
 	const r = 1
 
-	movparent := cur.children[l]
-	mov := movparent.children[r]
+	ln := cur.children[l]
+	lrn := ln.children[r]
 
-	mov.value, cur.value = cur.value, mov.value //交换值达到, 相对位移
+	lrln := lrn.children[l]
+	lrrn := lrn.children[r]
 
-	if mov.children[l] != nil {
-		movparent.children[r] = mov.children[l]
-		movparent.children[r].parent = movparent
-	} else {
-		movparent.children[r] = nil
-	}
+	tree.takeParent(cur, lrn)
 
-	if mov.children[r] != nil {
-		mov.children[l] = mov.children[r]
-	} else {
-		mov.children[l] = nil
-	}
+	setChild(ln, r, lrln)
+	setChild(cur, l, lrrn)
 
-	if cur.children[r] != nil {
-		mov.children[r] = cur.children[r]
-		mov.children[r].parent = mov
-	} else {
-		mov.children[r] = nil
-	}
+	setChildNotNil(lrn, l, ln)
+	setChildNotNil(lrn, r, cur)
 
-	cur.children[r] = mov
-	mov.parent = cur
-
-	movparent.size = getChildrenSumSize(movparent) + 1
-	mov.size = getChildrenSumSize(mov) + 1
+	ln.size = getChildrenSumSize(ln) + 1
 	cur.size = getChildrenSumSize(cur) + 1
-}
+	lrn.size = getChildrenSumSize(lrn) + 1
 
-func (tree *vbTree) replaceNotRoot(old, new *tNode) {
+	return lrn
 
-	new.children[0] = old.children[0]
-	new.children[1] = old.children[1]
+	// movparent := cur.children[l]
+	// mov := movparent.children[r]
 
-	if old.parent.children[1] == old {
-		old.parent.children[1] = new
-	} else {
-		old.parent.children[0] = new
-	}
+	// mov.value, cur.value = cur.value, mov.value //交换值达到, 相对位移
+
+	// if mov.children[l] != nil {
+	// 	movparent.children[r] = mov.children[l]
+	// 	movparent.children[r].parent = movparent
+	// } else {
+	// 	movparent.children[r] = nil
+	// }
+
+	// if mov.children[r] != nil {
+	// 	mov.children[l] = mov.children[r]
+	// } else {
+	// 	mov.children[l] = nil
+	// }
+
+	// if cur.children[r] != nil {
+	// 	mov.children[r] = cur.children[r]
+	// 	mov.children[r].parent = mov
+	// } else {
+	// 	mov.children[r] = nil
+	// }
+
+	// cur.children[r] = mov
+	// mov.parent = cur
+
+	// movparent.size = getChildrenSumSize(movparent) + 1
+	// mov.size = getChildrenSumSize(mov) + 1
+	// cur.size = getChildrenSumSize(cur) + 1
 }
 
 func (tree *vbTree) replace(old, new *tNode) {
@@ -794,60 +795,31 @@ func (tree *vbTree) rrotate3(cur *tNode) *tNode {
 	const r = 1
 	// 1 right 0 left
 	mov := cur.children[l]
-
-	if cur.parent == nil {
-		tree.root = mov
-	} else {
-		if cur.parent.children[1] == cur {
-			cur.parent.children[1] = mov
-		} else {
-			cur.parent.children[0] = mov
-		}
-	}
-	mov.parent = cur.parent
-
-	mov.children[r] = cur
-	mov.children[r].parent = mov
-
 	cur.children[l] = nil
 
+	tree.takeParent(cur, mov)
+	setChildNotNil(mov, r, cur)
+
 	mov.size = 3
 	cur.size = 1
 	return mov
 }
 
-func (tree *vbTree) rrotate(cur *tNode) {
-
+func (tree *vbTree) rrotate(cur *tNode) *tNode {
 	const l = 0
 	const r = 1
 	// 1 right 0 left
-	mov := cur.children[l]
+	ln := cur.children[l]
+	lrn := ln.children[r]
 
-	mov.value, cur.value = cur.value, mov.value //交换值达到, 相对位移
+	tree.takeParent(cur, ln)
+	setChild(cur, l, lrn)
+	setChildNotNil(ln, r, cur)
 
-	//  mov.children[l]不可能为nil
-	mov.children[l].parent = cur
-	cur.children[l] = mov.children[l]
-
-	// 解决mov节点孩子转移的问题
-	if mov.children[r] != nil {
-		mov.children[l] = mov.children[r]
-	} else {
-		mov.children[l] = nil
-	}
-
-	if cur.children[r] != nil {
-		mov.children[r] = cur.children[r]
-		mov.children[r].parent = mov
-	} else {
-		mov.children[r] = nil
-	}
-
-	// 连接转移后的节点 由于mov只是与cur交换值,parent不变
-	cur.children[r] = mov
-
-	mov.size = getChildrenSumSize(mov) + 1
 	cur.size = getChildrenSumSize(cur) + 1
+	ln.size = getChildrenSumSize(ln) + 1
+
+	return ln
 }
 
 func (tree *vbTree) lrotate3(cur *tNode) *tNode {
@@ -856,60 +828,33 @@ func (tree *vbTree) lrotate3(cur *tNode) *tNode {
 
 	// 1 right 0 left
 	mov := cur.children[l]
-
-	if cur.parent == nil {
-		tree.root = mov
-	} else {
-		if cur.parent.children[1] == cur {
-			cur.parent.children[1] = mov
-		} else {
-			cur.parent.children[0] = mov
-		}
-	}
-	mov.parent = cur.parent
-
-	mov.children[r] = cur
-	mov.children[r].parent = mov
-
 	cur.children[l] = nil
 
+	tree.takeParent(cur, mov)
+	setChildNotNil(mov, r, cur)
+
 	mov.size = 3
 	cur.size = 1
 	return mov
 }
 
-func (tree *vbTree) lrotate(cur *tNode) {
+func (tree *vbTree) lrotate(cur *tNode) *tNode {
 
 	const l = 1
 	const r = 0
+
 	// 1 right 0 left
-	mov := cur.children[l]
+	ln := cur.children[l]
+	lrn := ln.children[r]
 
-	mov.value, cur.value = cur.value, mov.value //交换值达到, 相对位移
+	tree.takeParent(cur, ln)
+	setChild(cur, l, lrn)
+	setChildNotNil(ln, r, cur)
 
-	//  mov.children[l]不可能为nil
-	mov.children[l].parent = cur
-	cur.children[l] = mov.children[l]
-
-	// 解决mov节点孩子转移的问题
-	if mov.children[r] != nil {
-		mov.children[l] = mov.children[r]
-	} else {
-		mov.children[l] = nil
-	}
-
-	if cur.children[r] != nil {
-		mov.children[r] = cur.children[r]
-		mov.children[r].parent = mov
-	} else {
-		mov.children[r] = nil
-	}
-
-	// 连接转移后的节点 由于mov只是与cur交换值,parent不变
-	cur.children[r] = mov
-
-	mov.size = getChildrenSumSize(mov) + 1
 	cur.size = getChildrenSumSize(cur) + 1
+	ln.size = getChildrenSumSize(ln) + 1
+
+	return ln
 }
 
 func getChildrenSumSize(cur *tNode) int {
@@ -933,28 +878,27 @@ func (tree *vbTree) fixSizeWithRemove(cur *tNode) {
 		if cur.size > 8 {
 			factor := cur.size / 10 // or factor = 1
 			if cur.children[1].size >= cur.children[0].size*2+factor || cur.children[0].size >= cur.children[1].size*2+factor {
-				tree.fixSize(cur)
+				cur = tree.fixSize(cur)
 			}
 		}
 		cur = cur.parent
 	}
 }
 
-func (tree *vbTree) fixSize(cur *tNode) {
+func (tree *vbTree) fixSize(cur *tNode) *tNode {
 	if cur.children[0].size > cur.children[1].size {
 		llsize, lrsize := getChildrenSize(cur.children[0])
 		if lrsize > llsize {
-			tree.rlrotate(cur)
-		} else {
-			tree.rrotate(cur)
+			return tree.rlrotate(cur)
 		}
+		return tree.rrotate(cur)
+
 	} else {
 		rlsize, rrsize := getChildrenSize(cur.children[1])
 		if rlsize > rrsize {
-			tree.lrrotate(cur)
-		} else {
-			tree.lrotate(cur)
+			return tree.lrrotate(cur)
 		}
+		return tree.lrotate(cur)
 	}
 }
 
diff --git a/priority_queue/vbt_test.go b/priority_queue/vbt_test.go
index d343ac2..0b4501f 100644
--- a/priority_queue/vbt_test.go
+++ b/priority_queue/vbt_test.go
@@ -330,14 +330,14 @@ func TestTravalsal(t *testing.T) {
 
 func TestRemoveAll(t *testing.T) {
 ALL:
-	for c := 0; c < 5000; c++ {
+	for c := 0; c < 10000; c++ {
 		tree := newVBT(compare.Int)
 		gods := avltree.NewWithIntComparator()
 		var l []int
 		m := make(map[int]int)
 
-		for i := 0; len(l) < 10; i++ {
-			v := randomdata.Number(0, 100)
+		for i := 0; len(l) < 50; i++ {
+			v := randomdata.Number(0, 65535)
 			if _, ok := m[v]; !ok {
 				m[v] = v
 				l = append(l, v)
@@ -346,14 +346,10 @@ ALL:
 			}
 		}
 
-		for i := 0; i < 10; i++ {
-			beforce := tree.debugString()
+		for i := 0; i < 50; i++ {
 			tree.Remove(l[i])
-			after := tree.debugString()
 			gods.Remove(l[i])
 
-			t.Error(beforce)
-			t.Error(after, l[i])
 			s1 := spew.Sprint(tree.Values())
 			s2 := spew.Sprint(gods.Values())
 			if s1 != s2 {
@@ -628,7 +624,7 @@ func BenchmarkPut(b *testing.B) {
 	b.ResetTimer()
 	b.StartTimer()
 
-	execCount := 10
+	execCount := 50
 	b.N = len(l) * execCount
 	for i := 0; i < execCount; i++ {
 		tree := newVBT(compare.Int)