From 15d2f6acd3779b31022a8f80256eafd217db4ae6 Mon Sep 17 00:00:00 2001
From: eson <474420502@qq.com>
Date: Fri, 29 Mar 2019 02:44:24 +0800
Subject: [PATCH] =?UTF-8?q?removeNode=20=E6=9C=89=E9=94=99?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 priority_queue/priority_queue.go      | 14 +++++-
 priority_queue/priority_queue_test.go | 71 +++++++++++++++++++++++++++
 priority_queue/vbt.go                 |  2 +-
 3 files changed, 84 insertions(+), 3 deletions(-)
 create mode 100644 priority_queue/priority_queue_test.go

diff --git a/priority_queue/priority_queue.go b/priority_queue/priority_queue.go
index 554a1e4..b56476e 100644
--- a/priority_queue/priority_queue.go
+++ b/priority_queue/priority_queue.go
@@ -26,8 +26,7 @@ func (pq *PriorityQueue) Push(value interface{}) {
 	if pq.head == nil {
 		pq.head = n
 		return
-	}
-	if pq.queue.Compare(n, pq.head) == 1 {
+	} else if pq.queue.Compare(n.value, pq.head.value) == 1 {
 		pq.head = n
 	}
 }
@@ -40,6 +39,17 @@ func (pq *PriorityQueue) Top() (result interface{}, ok bool) {
 }
 
 func (pq *PriorityQueue) Pop() (result interface{}, ok bool) {
+	if pq.head != nil {
+		prev := getPrev(pq.head, 1)
+		result = pq.head.value
+		pq.queue.removeNode(pq.head)
+		if prev != nil {
+			pq.head = prev
+		} else {
+			pq.head = nil
+		}
+		return result, true
+	}
 	return nil, false
 }
 
diff --git a/priority_queue/priority_queue_test.go b/priority_queue/priority_queue_test.go
new file mode 100644
index 0000000..c65521c
--- /dev/null
+++ b/priority_queue/priority_queue_test.go
@@ -0,0 +1,71 @@
+package pqueue
+
+import (
+	"testing"
+
+	"github.com/Pallinder/go-randomdata"
+
+	"474420502.top/eson/structure/compare"
+)
+
+func TestPush(t *testing.T) {
+	pq := New(compare.Int)
+	for i := 0; i < 10; i++ {
+		v := randomdata.Number(0, 100)
+		pq.Push(v)
+		t.Error(v)
+	}
+	t.Error(pq.Pop())
+	t.Error(pq.Pop())
+	t.Error(pq.Pop())
+	t.Error(pq.Pop())
+	t.Error(pq.Pop())
+}
+
+func TestPop(t *testing.T) {
+	pq := New(compare.Int)
+	for i := 0; i < 10; i++ {
+		v := randomdata.Number(0, 100)
+		pq.Push(v)
+		t.Error(v)
+	}
+
+	for i := 0; i < 10; i++ {
+		t.Error(pq.Pop())
+	}
+}
+
+func BenchmarkPriorityPush(b *testing.B) {
+
+	l := loadTestData()
+	execCount := 5
+	b.N = len(l) * execCount
+
+	b.ResetTimer()
+	b.StartTimer()
+
+	for i := 0; i < execCount; i++ {
+		pq := New(compare.Int)
+		for _, v := range l {
+			pq.Push(v)
+		}
+	}
+}
+
+func BenchmarkPriorityPop(b *testing.B) {
+
+	l := loadTestData()
+
+	pq := New(compare.Int)
+	for _, v := range l {
+		pq.Push(v)
+	}
+
+	b.N = len(l) / 1000
+	b.ResetTimer()
+	b.StartTimer()
+
+	for i := 0; i < b.N; i++ {
+		pq.Pop()
+	}
+}
diff --git a/priority_queue/vbt.go b/priority_queue/vbt.go
index cb4e7de..1f584a2 100644
--- a/priority_queue/vbt.go
+++ b/priority_queue/vbt.go
@@ -889,7 +889,7 @@ func (tree *vbTree) fixSizeWithRemove(cur *tNode) {
 			ls, rs := getChildrenSize(cur)
 			factor := cur.size / 10 // or factor = 1
 			if rs >= ls*2+factor || ls >= rs*2+factor {
-				tree.fixSize(cur, ls, rs)
+				cur = tree.fixSize(cur, ls, rs)
 			}
 		}
 		cur = cur.parent