diff --git a/max_heap.go b/max_heap.go index 888bdce..28b57fb 100644 --- a/max_heap.go +++ b/max_heap.go @@ -83,6 +83,16 @@ func (h *MaxHeap) Items() []*Item { return h.items } +// Clone returns a copy of the queue which can be modified without affecting +// the original queue. +func (h *MaxHeap) Clone() *MaxHeap { + c := NewMaxHeap(len(h.items)) + for _, item := range h.items { + c.items = append(c.items, &Item{ID: item.ID, Prio: item.Prio}) + } + return c +} + // PopAndPush removes the item with the highest priority value and adds a new // value to the heap in one operation. This is faster than two separate calls // to Pop and Push. diff --git a/min_heap.go b/min_heap.go index 6644c97..a312853 100644 --- a/min_heap.go +++ b/min_heap.go @@ -87,6 +87,16 @@ func (h *MinHeap) Items() []*Item { return h.items } +// Clone returns a copy of the queue which can be modified without affecting +// the original queue. +func (h *MinHeap) Clone() *MinHeap { + c := NewMinHeap(len(h.items)) + for _, item := range h.items { + c.items = append(c.items, &Item{ID: item.ID, Prio: item.Prio}) + } + return c +} + // PopAndPush removes the item with the lowest priority value and adds a new // value to the heap in one operation. This is faster than two separate calls // to Pop and Push. diff --git a/topk/topk.go b/topk/topk.go index 720173f..9bbe199 100644 --- a/topk/topk.go +++ b/topk/topk.go @@ -1,5 +1,6 @@ // Package topk selects the items with the lowest or highest priority from a -// larger set of items using a bounded priority queue. +// larger set of items using a bounded priority queue. This needs O(n log k) +// time and O(k) memory instead of sorting all n items. package topk import "github.com/fgrosse/prioqueue" diff --git a/topk/topk_test.go b/topk/topk_test.go index 4f81cf1..72637c1 100644 --- a/topk/topk_test.go +++ b/topk/topk_test.go @@ -19,3 +19,13 @@ func TestSmallest(t *testing.T) { expected := []prioqueue.Item{{ID: 2, Prio: 1}, {ID: 4, Prio: 3}} assert.Equal(t, expected, Smallest(items, 2)) } + +func TestLargest(t *testing.T) { + expected := []prioqueue.Item{{ID: 5, Prio: 9}, {ID: 3, Prio: 8}, {ID: 1, Prio: 5}} + assert.Equal(t, expected, Largest(items, 3)) +} + +func TestInvalidK(t *testing.T) { + assert.Nil(t, Smallest(items, 0)) + assert.Nil(t, Largest(items, -1)) +}