diff --git a/kpq/keyed_priority_queue.go b/kpq/keyed_priority_queue.go index 1627c45..bf98021 100644 --- a/kpq/keyed_priority_queue.go +++ b/kpq/keyed_priority_queue.go @@ -229,14 +229,14 @@ func (pq *KeyedPriorityQueue[K, V]) ValueOf(k K) (V, bool) { } // Remove removes the priority value associated with the given key k from the priority queue. -// It's a no-op if there's no such key k in the priority queue. -func (pq *KeyedPriorityQueue[K, V]) Remove(k K) { +// It's a no-op if there's no such key k in the priority queue and it will return false. +func (pq *KeyedPriorityQueue[K, V]) Remove(k K) bool { pq.mu.Lock() defer pq.mu.Unlock() i, ok := pq.im[k] if !ok { - return + return false } n := len(pq.pm) - 1 if i != n { @@ -247,6 +247,7 @@ func (pq *KeyedPriorityQueue[K, V]) Remove(k K) { pq.pm = pq.pm[:n] delete(pq.im, k) delete(pq.vals, k) + return true } // Len returns the size of the priority queue. diff --git a/kpq/keyed_priority_queue_test.go b/kpq/keyed_priority_queue_test.go index 9b1b91d..100ade1 100644 --- a/kpq/keyed_priority_queue_test.go +++ b/kpq/keyed_priority_queue_test.go @@ -286,39 +286,60 @@ func TestKeyedPriorityQueue_Remove(t *testing.T) { t.Run("Keys", func(t *testing.T) { testCases := []struct { key string + wantStatus bool wantPeekKey string wantPeekValue int wantLen int }{ { key: "first", + wantStatus: true, wantPeekKey: "second", wantPeekValue: 8, wantLen: 4, }, { key: "third", + wantStatus: true, wantPeekKey: "second", wantPeekValue: 8, wantLen: 3, }, { key: "second", + wantStatus: true, wantPeekKey: "fourth", wantPeekValue: 10, wantLen: 2, }, { key: "last", + wantStatus: true, wantPeekKey: "fourth", wantPeekValue: 10, wantLen: 1, }, + { + key: "nonexistent", + wantStatus: false, + wantLen: 1, + }, } for _, tc := range testCases { t.Run(tc.key, func(t *testing.T) { - pq.Remove(tc.key) + gotStatus := pq.Remove(tc.key) + if gotStatus != tc.wantStatus { + t.Errorf("pq.Remove(): got status %t; want %t", gotStatus, tc.wantStatus) + } + + if got := pq.Len(); got != tc.wantLen { + t.Errorf("pq.Len(): got %d; want %d", got, tc.wantLen) + } + + if !gotStatus { + return + } gotPeekKey, gotPeekValue, ok := pq.Peek() if !ok { @@ -332,22 +353,9 @@ func TestKeyedPriorityQueue_Remove(t *testing.T) { if gotPeekValue != tc.wantPeekValue { t.Errorf("pq.PeekValue(): got value %d; want %d", gotPeekValue, tc.wantPeekValue) } - - if got := pq.Len(); got != tc.wantLen { - t.Errorf("pq.Len(): got %d; want %d", got, tc.wantLen) - } }) } }) - - t.Run("NonExistingKey", func(t *testing.T) { - want := pq.Len() - pq.Remove("non-existing-key") - - if got := pq.Len(); got != want { - t.Errorf("pq.Len(): got %d; want %d", got, want) - } - }) } func TestKeyedPriorityQueue_Peek_EmptyQueue(t *testing.T) { @@ -479,7 +487,6 @@ func TestKeyedPriorityQueue_Set(t *testing.T) { } }) } - } func benchmarkKeyedPriorityQueue_PushPop(b *testing.B, n int) { @@ -500,18 +507,23 @@ func benchmarkKeyedPriorityQueue_PushPop(b *testing.B, n int) { func BenchmarkKeyedPriorityQueue_PushPop_10(b *testing.B) { benchmarkKeyedPriorityQueue_PushPop(b, 10) } + func BenchmarkKeyedPriorityQueue_PushPop_100(b *testing.B) { benchmarkKeyedPriorityQueue_PushPop(b, 100) } + func BenchmarkKeyedPriorityQueue_PushPop_1000(b *testing.B) { benchmarkKeyedPriorityQueue_PushPop(b, 1000) } + func BenchmarkKeyedPriorityQueue_PushPop_10000(b *testing.B) { benchmarkKeyedPriorityQueue_PushPop(b, 10000) } + func BenchmarkKeyedPriorityQueue_PushPop_100000(b *testing.B) { benchmarkKeyedPriorityQueue_PushPop(b, 100000) } + func BenchmarkKeyedPriorityQueue_PushPop_1000000(b *testing.B) { benchmarkKeyedPriorityQueue_PushPop(b, 1000000) }