diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..fa45d85 --- /dev/null +++ b/go.mod @@ -0,0 +1,3 @@ +module github.com/airspace-link-inc/go-priority-queue + +go 1.18 diff --git a/priorty_queue.go b/priorty_queue.go index 7541dd9..9f93200 100644 --- a/priorty_queue.go +++ b/priorty_queue.go @@ -11,33 +11,34 @@ import ( "errors" ) -// PriorityQueue represents the queue -type PriorityQueue struct { - itemHeap *itemHeap - lookup map[interface{}]*item +var ErrEmptyQueue = errors.New("empty queue") + +// PriorityQueue[T] represents the queue of comparable type T. +type PriorityQueue[T comparable] struct { + itemHeap *itemHeap[T] + lookup map[T]*item[T] } -// New initializes an empty priority queue. -func New() PriorityQueue { - return PriorityQueue{ - itemHeap: &itemHeap{}, - lookup: make(map[interface{}]*item), +// New[T] initializes an empty priority queue of type T. +func New[T comparable]() PriorityQueue[T] { + return PriorityQueue[T]{ + itemHeap: &itemHeap[T]{}, + lookup: make(map[T]*item[T]), } } // Len returns the number of elements in the queue. -func (p *PriorityQueue) Len() int { +func (p *PriorityQueue[T]) Len() int { return p.itemHeap.Len() } // Insert inserts a new element into the queue. No action is performed on duplicate elements. -func (p *PriorityQueue) Insert(v interface{}, priority float64) { - _, ok := p.lookup[v] - if ok { +func (p *PriorityQueue[T]) Insert(v T, priority float64) { + if _, ok := p.lookup[v]; ok { return } - newItem := &item{ + newItem := &item[T]{ value: v, priority: priority, } @@ -47,19 +48,20 @@ func (p *PriorityQueue) Insert(v interface{}, priority float64) { // Pop removes the element with the highest priority from the queue and returns it. // In case of an empty queue, an error is returned. -func (p *PriorityQueue) Pop() (interface{}, error) { +func (p *PriorityQueue[T]) Pop() (T, error) { if len(*p.itemHeap) == 0 { - return nil, errors.New("empty queue") + var zeroVal T + return zeroVal, ErrEmptyQueue } - item := heap.Pop(p.itemHeap).(*item) + item := heap.Pop(p.itemHeap).(*item[T]) delete(p.lookup, item.value) return item.value, nil } // UpdatePriority changes the priority of a given item. // If the specified item is not present in the queue, no action is performed. -func (p *PriorityQueue) UpdatePriority(x interface{}, newPriority float64) { +func (p *PriorityQueue[T]) UpdatePriority(x T, newPriority float64) { item, ok := p.lookup[x] if !ok { return @@ -69,35 +71,35 @@ func (p *PriorityQueue) UpdatePriority(x interface{}, newPriority float64) { heap.Fix(p.itemHeap, item.index) } -type itemHeap []*item +type itemHeap[T comparable] []*item[T] -type item struct { - value interface{} +type item[T comparable] struct { + value T priority float64 index int } -func (ih *itemHeap) Len() int { +func (ih *itemHeap[T]) Len() int { return len(*ih) } -func (ih *itemHeap) Less(i, j int) bool { +func (ih *itemHeap[T]) Less(i, j int) bool { return (*ih)[i].priority < (*ih)[j].priority } -func (ih *itemHeap) Swap(i, j int) { +func (ih *itemHeap[T]) Swap(i, j int) { (*ih)[i], (*ih)[j] = (*ih)[j], (*ih)[i] (*ih)[i].index = i (*ih)[j].index = j } -func (ih *itemHeap) Push(x interface{}) { - it := x.(*item) +func (ih *itemHeap[T]) Push(x interface{}) { + it := x.(*item[T]) it.index = len(*ih) *ih = append(*ih, it) } -func (ih *itemHeap) Pop() interface{} { +func (ih *itemHeap[T]) Pop() interface{} { old := *ih item := old[len(old)-1] *ih = old[0 : len(old)-1] diff --git a/priorty_queue_test.go b/priorty_queue_test.go index be6d04f..710e8ed 100644 --- a/priorty_queue_test.go +++ b/priorty_queue_test.go @@ -1,12 +1,13 @@ package pq import ( + "math" "sort" "testing" ) func TestPriorityQueue(t *testing.T) { - pq := New() + pq := New[float64]() elements := []float64{5, 3, 7, 8, 6, 2, 9} for _, e := range elements { pq.Insert(e, e) @@ -19,7 +20,7 @@ func TestPriorityQueue(t *testing.T) { t.Fatalf(err.Error()) } - i := item.(float64) + i := item if e != i { t.Fatalf("expected %v, got %v", e, i) } @@ -27,7 +28,7 @@ func TestPriorityQueue(t *testing.T) { } func TestPriorityQueueUpdate(t *testing.T) { - pq := New() + pq := New[string]() pq.Insert("foo", 3) pq.Insert("bar", 4) pq.UpdatePriority("bar", 2) @@ -37,13 +38,13 @@ func TestPriorityQueueUpdate(t *testing.T) { t.Fatal(err.Error()) } - if item.(string) != "bar" { + if item != "bar" { t.Fatal("priority update failed") } } func TestPriorityQueueLen(t *testing.T) { - pq := New() + pq := New[string]() if pq.Len() != 0 { t.Fatal("empty queue should have length of 0") } @@ -55,8 +56,24 @@ func TestPriorityQueueLen(t *testing.T) { } } +func TestItemHeapLess(t *testing.T) { + h := itemHeap[int]{ + &item[int]{priority: math.Inf(1)}, + &item[int]{priority: math.Inf(1)}, + &item[int]{priority: math.Inf(-1)}, + &item[int]{priority: math.Inf(-1)}, + } + + // test all pairwise elements: 0,1 then 2,3 then... + for i := 0; i < len(h); i += 2 { + if h.Less(i, i+1) { + t.Fatalf("%v should not have less priority than %v", h[i], h[i+1]) + } + } +} + func TestDoubleAddition(t *testing.T) { - pq := New() + pq := New[string]() pq.Insert("foo", 2) pq.Insert("bar", 3) pq.Insert("bar", 1) @@ -66,21 +83,23 @@ func TestDoubleAddition(t *testing.T) { } item, _ := pq.Pop() - if item.(string) != "foo" { + if item != "foo" { t.Fatal("queue should ignore duplicate insert, not update existing item") } } func TestPopEmptyQueue(t *testing.T) { - pq := New() + pq := New[float32]() _, err := pq.Pop() if err == nil { t.Fatal("should produce error when performing pop on empty queue") + } else if err != ErrEmptyQueue { + t.Fatalf("error should be equal to %v, got %v", ErrEmptyQueue, err) } } func TestUpdateNonExistingItem(t *testing.T) { - pq := New() + pq := New[string]() pq.Insert("foo", 4) pq.UpdatePriority("bar", 5) @@ -90,7 +109,7 @@ func TestUpdateNonExistingItem(t *testing.T) { } item, _ := pq.Pop() - if item.(string) != "foo" { - t.Fatalf("update should not overwrite item, expected \"foo\", got \"%v\"", item.(string)) + if item != "foo" { + t.Fatalf("update should not overwrite item, expected \"foo\", got \"%v\"", item) } }