From 17f800dd3e95cd44957261b3ff3d0da039ffd303 Mon Sep 17 00:00:00 2001
From: Dirk McCormick <dirkmdev@gmail.com>
Date: Fri, 1 Nov 2019 10:59:20 -0400
Subject: [PATCH] feat: Remove()

---
 pq.go      |  9 ++++++---
 pq_test.go | 22 ++++++++++++++++++++++
 2 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/pq.go b/pq.go
index 6e5522e..10eaf7f 100644
--- a/pq.go
+++ b/pq.go
@@ -11,13 +11,12 @@ type PQ interface {
 	Pop() Elem
 	// Peek returns the highest priority Elem in PQ (without removing it).
 	Peek() Elem
+	// Remove removes the item at the given index from the PQ.
+	Remove(index int) Elem
 	// Len returns the number of elements in the PQ.
 	Len() int
 	// Update `fixes` the PQ.
 	Update(index int)
-
-	// TODO  explain why this interface should not be extended
-	// It does not support Remove. This is because...
 }
 
 // Elem describes elements that can be added to the PQ. Clients must implement
@@ -65,6 +64,10 @@ func (w *wrapper) Peek() Elem {
 	return w.heapinterface.elems[0].(Elem)
 }
 
+func (w *wrapper) Remove(index int) Elem {
+	return heap.Remove(&w.heapinterface, index).(Elem)
+}
+
 func (w *wrapper) Update(index int) {
 	heap.Fix(&w.heapinterface, index)
 }
diff --git a/pq_test.go b/pq_test.go
index f63911c..5e53628 100644
--- a/pq_test.go
+++ b/pq_test.go
@@ -67,6 +67,28 @@ func TestCorrectnessOfPop(t *testing.T) {
 	}
 }
 
+func TestRemove(t *testing.T) {
+	q := New(PriorityComparator)
+	tasks := []TestElem{
+		{Key: "a", Priority: 9},
+		{Key: "b", Priority: 4},
+		{Key: "c", Priority: 3},
+	}
+	for i := range tasks {
+		q.Push(&tasks[i])
+	}
+	removed := q.Remove(1).(*TestElem)
+	if q.Len() != 2 {
+		t.Fatal("expected item to have been removed")
+	}
+	if q.Pop().(*TestElem).Key == removed.Key {
+		t.Fatal("Remove() returned wrong element")
+	}
+	if q.Pop().(*TestElem).Key == removed.Key {
+		t.Fatal("Remove() returned wrong element")
+	}
+}
+
 func TestUpdate(t *testing.T) {
 	t.Log(`
 	Add 3 elements.
-- 
GitLab