Remove()

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)
 }

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.