diff --git a/pq.go b/pq.go index 6e5522e65aa884545ce037be705e2e450003a550..10eaf7f8862df1e10f32ccf558b679cd3af13673 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 f63911c14de753c088565d6ceb8668fb10ccdcb7..5e53628804b61b53b181e1987088a2c6925ff4fa 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.