Repository: dominikbraun/graph
Pull request #43 (2022-09-12 19:49:22)

Priority Queue implementation using heap
Hi @dominikbraun, this PR is for issue #29.
This implementation is mostly from this [Priority Queue example](https://pkg.go.dev/container/heap#example-package-PriorityQueue) from `container/heap` package.

Let me know if something needed to be updated since I made some adjustments on `collection.go` and its tests.
Thank you!


Patch:
diff --git a/collection.go b/collection.go
index b166241..a220c4b 100644
--- a/collection.go
+++ b/collection.go
@@ -1,94 +1,111 @@
 package graph
 
 import (
+	"container/heap"
 	"errors"
-	"sort"
 )
 
 // priorityQueue is a priority queue implementation for minimum priorities, meaning that smaller
 // values will be prioritized. It maintains a descendingly ordered list of priority items.
 //
-// This is still a naive implementation, which is to be replaced with a binary heap implementation.
+// This implementation is built on top of heap.Interface with some adjustments to comply with our
+// generic usage.
 type priorityQueue[T comparable] struct {
-	items []priorityItem[T]
+	items *priorityItems[T]
+
+	// The map used to look up item that already pushed to the queue. Especially useful when we do
+	// UpdatePriority operation.
+	lookup map[T]*priorityItem[T]
 }
 
+type priorityItems[T comparable] []*priorityItem[T]
+
 // priorityItem is an item in the priority queue, consisting of a priority and an actual value.
 type priorityItem[T comparable] struct {
 	value    T
 	priority float64
+
+	// The index field is used and operated internally by heap.Interface to re-organize items in the
+	// queue.
+	index int
 }
 
 func newPriorityQueue[T comparable]() *priorityQueue[T] {
 	return &priorityQueue[T]{
-		items: make([]priorityItem[T], 0),
+		items:  &priorityItems[T]{},
+		lookup: map[T]*priorityItem[T]{},
 	}
 }
 
-// Push pushes a new item with the given priority into the queue. Because Push keeps track of the
-// descendant order of items and priorities, it has O(n) insertion time.
+func (p *priorityQueue[T]) Len() int {
+	return p.items.Len()
+}
+
+// Push pushes a new item with the given priority into the queue.
 func (p *priorityQueue[T]) Push(item T, priority float64) {
-	index := p.Len() - 1
-
-	for i := p.Len(); i > 0; i-- {
-		currentItem := p.items[i-1]
-		if currentItem.priority > priority {
-			index = i
-			break
-		}
+	if _, ok := p.lookup[item]; ok {
+		return
 	}
 
-	p.insertItemAt(item, priority, index)
+	newItem := &priorityItem[T]{
+		value:    item,
+		priority: priority,
+		index:    0,
+	}
+
+	heap.Push(p.items, newItem)
+	p.lookup[item] = newItem
 }
 
-// Pop returns the item with the smallest priority from the queue and removes that item. Returns an
-// error if the priority queue is empty, which can be tested using Len first.s
+// Pop returns the item with the smallest priority from the queue and removes that item.
 func (p *priorityQueue[T]) Pop() (T, error) {
-	var priorityItem priorityItem[T]
-
-	if p.Len() == 0 {
-		return priorityItem.value, errors.New("priority queue is empty")
+	if len(*p.items) == 0 {
+		var zeroVal T
+		return zeroVal, errors.New("priority queue is empty")
 	}
 
-	priorityItem = p.items[p.Len()-1]
-	p.items = p.items[:p.Len()-1]
-
-	return priorityItem.value, nil
+	item := heap.Pop(p.items).(*priorityItem[T])
+	delete(p.lookup, item.value)
+	return item.value, nil
 }
 
-// DecreasePriority decreases the priority of a given item to the given priority. The item must be
+// UpdatePriority updates the priority of a given item to the given priority. The item must be
 // pushed into the queue first. If the item doesn't exist, nothing happens.
-//
-// With the current implementation, DecreasePriority causes the items in the queue to be re-sorted.
-func (p *priorityQueue[T]) DecreasePriority(item T, priority float64) {
-	for i, currentItem := range p.items {
-		if currentItem.value == item {
-			p.items[i].priority = priority
-		}
+func (p *priorityQueue[T]) UpdatePriority(item T, priority float64) {
+	targetItem, ok := p.lookup[item]
+	if !ok {
+		return
 	}
 
-	sort.Slice(p.items, func(i, j int) bool {
-		return p.items[i].priority > p.items[j].priority
-	})
+	targetItem.priority = priority
+	heap.Fix(p.items, targetItem.index)
 }
 
-// Len returns the current length of the priority queue, i.e. the number of items in the queue.
-func (p *priorityQueue[T]) Len() int {
-	return len(p.items)
+// Len and the rest methods below are required to implement priority queue on top of heap.Interface.
+func (pi *priorityItems[T]) Len() int {
+	return len(*pi)
 }
 
-func (p *priorityQueue[T]) insertItemAt(item T, priority float64, index int) {
-	priorityItem := priorityItem[T]{
-		value:    item,
-		priority: priority,
-	}
+func (pi *priorityItems[T]) Less(i, j int) bool {
+	return (*pi)[i].priority < (*pi)[j].priority
+}
 
-	if p.Len() == 0 || p.Len() == index {
-		p.items = append(p.items, priorityItem)
-		return
-	}
+func (pi *priorityItems[T]) Swap(i, j int) {
+	(*pi)[i], (*pi)[j] = (*pi)[j], (*pi)[i]
+	(*pi)[i].index = i
+	(*pi)[j].index = j
+}
+
+func (pi *priorityItems[T]) Push(x interface{}) {
+	it := x.(*priorityItem[T])
+	it.index = len(*pi)
+	*pi = append(*pi, it)
+}
 
-	p.items = append(p.items[:index+1], p.items[index:]...)
+func (pi *priorityItems[T]) Pop() interface{} {
+	old := *pi
+	item := old[len(old)-1]
+	*pi = old[:len(old)-1]
 
-	p.items[index] = priorityItem
+	return item
 }
diff --git a/paths.go b/paths.go
index d8d5ba6..76eeca0 100644
--- a/paths.go
+++ b/paths.go
@@ -118,7 +118,7 @@ func ShortestPath[K comparable, T any](g Graph[K, T], source, target K) ([]K, er
 			if weight < weights[adjacency] && !hasInfiniteWeight {
 				weights[adjacency] = weight
 				bestPredecessors[adjacency] = vertex
-				queue.DecreasePriority(adjacency, weight)
+				queue.UpdatePriority(adjacency, weight)
 			}
 		}
 	}


Test patch:
diff --git a/collection_test.go b/collection_test.go
index 494c013..089e071 100644
--- a/collection_test.go
+++ b/collection_test.go
@@ -4,16 +4,38 @@ import (
 	"testing"
 )
 
+func assertPriorityQueue[T comparable](t *testing.T, testName string, q *priorityQueue[T], expectedItems []*priorityItem[T]) {
+	if q.Len() != len(expectedItems) {
+		t.Fatalf("%s: item length expectancy doesn't match: expected %v, got %v", testName, len(expectedItems), q.Len())
+	}
+
+	poppedList := make([]T, q.Len())
+	for q.Len() > 0 {
+		poppedItem, _ := q.Pop()
+		poppedList = append(poppedList, poppedItem)
+	}
+
+	// Each time we pop from the given queue, it will return the item with the smallest priority.
+	// So we compare the expected item with popped list from the other direction.
+	n := len(poppedList)
+	for i, expectedItem := range expectedItems {
+		comparedToItem := poppedList[n-1-i]
+		if comparedToItem != expectedItem.value {
+			t.Errorf("%s: item doesn't match: expected %v at index %d, got %v", testName, expectedItem.value, i, comparedToItem)
+		}
+	}
+}
+
 func TestPriorityQueue_Push(t *testing.T) {
 	tests := map[string]struct {
 		items                 []int
 		priorities            []float64
-		expectedPriorityItems []priorityItem[int]
+		expectedPriorityItems []*priorityItem[int]
 	}{
 		"queue with 5 elements": {
 			items:      []int{10, 20, 30, 40, 50},
 			priorities: []float64{6, 8, 2, 7, 5},
-			expectedPriorityItems: []priorityItem[int]{
+			expectedPriorityItems: []*priorityItem[int]{
 				{value: 20, priority: 8},
 				{value: 40, priority: 7},
 				{value: 10, priority: 6},
@@ -30,15 +52,7 @@ func TestPriorityQueue_Push(t *testing.T) {
 			queue.Push(item, test.priorities[i])
 		}
 
-		if len(queue.items) != len(test.expectedPriorityItems) {
-			t.Fatalf("%s: item length expectancy doesn't match: expected %v, got %v", name, len(test.expectedPriorityItems), len(queue.items))
-		}
-
-		for i, expectedPriorityItem := range test.expectedPriorityItems {
-			if queue.items[i] != expectedPriorityItem {
-				t.Errorf("%s: item doesn't match: expected %v at index %d, got %v", name, expectedPriorityItem, i, queue.items[i])
-			}
-		}
+		assertPriorityQueue(t, name, queue, test.expectedPriorityItems)
 	}
 }
 
@@ -87,15 +101,15 @@ func TestPriorityQueue_Pop(t *testing.T) {
 	}
 }
 
-func TestPriorityQueue_DecreasePriority(t *testing.T) {
+func TestPriorityQueue_UpdatePriority(t *testing.T) {
 	tests := map[string]struct {
-		items                 []priorityItem[int]
-		expectedPriorityItems []priorityItem[int]
+		items                 []*priorityItem[int]
+		expectedPriorityItems []*priorityItem[int]
 		decreaseItem          int
 		decreasePriority      float64
 	}{
 		"decrease 30 to priority 5": {
-			items: []priorityItem[int]{
+			items: []*priorityItem[int]{
 				{value: 40, priority: 40},
 				{value: 30, priority: 30},
 				{value: 20, priority: 20},
@@ -103,7 +117,7 @@ func TestPriorityQueue_DecreasePriority(t *testing.T) {
 			},
 			decreaseItem:     30,
 			decreasePriority: 5,
-			expectedPriorityItems: []priorityItem[int]{
+			expectedPriorityItems: []*priorityItem[int]{
 				{value: 40, priority: 40},
 				{value: 20, priority: 20},
 				{value: 10, priority: 10},
@@ -111,7 +125,7 @@ func TestPriorityQueue_DecreasePriority(t *testing.T) {
 			},
 		},
 		"decrease a non-existent item": {
-			items: []priorityItem[int]{
+			items: []*priorityItem[int]{
 				{value: 40, priority: 40},
 				{value: 30, priority: 30},
 				{value: 20, priority: 20},
@@ -119,31 +133,41 @@ func TestPriorityQueue_DecreasePriority(t *testing.T) {
 			},
 			decreaseItem:     50,
 			decreasePriority: 10,
-			expectedPriorityItems: []priorityItem[int]{
+			expectedPriorityItems: []*priorityItem[int]{
 				{value: 40, priority: 40},
 				{value: 30, priority: 30},
 				{value: 20, priority: 20},
 				{value: 10, priority: 10},
 			},
 		},
+		"increase 10 to priority 100": {
+			items: []*priorityItem[int]{
+				{value: 40, priority: 40},
+				{value: 30, priority: 30},
+				{value: 20, priority: 20},
+				{value: 10, priority: 10},
+			},
+			decreaseItem:     10,
+			decreasePriority: 100,
+			expectedPriorityItems: []*priorityItem[int]{
+				{value: 10, priority: 100},
+				{value: 40, priority: 40},
+				{value: 30, priority: 30},
+				{value: 20, priority: 20},
+			},
+		},
 	}
 
 	for name, test := range tests {
-		queue := &priorityQueue[int]{
-			items: test.items,
-		}
-
-		queue.DecreasePriority(test.decreaseItem, test.decreasePriority)
+		queue := newPriorityQueue[int]()
 
-		if len(queue.items) != len(test.expectedPriorityItems) {
-			t.Fatalf("%s: item length expectancy doesn't match: expected %v, got %v", name, len(test.expectedPriorityItems), len(queue.items))
+		for _, item := range test.items {
+			queue.Push(item.value, item.priority)
 		}
 
-		for i, expectedPriorityItem := range test.expectedPriorityItems {
-			if queue.items[i] != expectedPriorityItem {
-				t.Errorf("%s: item doesn't match: expected %v at index %d, got %v", name, expectedPriorityItem, i, queue.items[i])
-			}
-		}
+		queue.UpdatePriority(test.decreaseItem, test.decreasePriority)
+
+		assertPriorityQueue(t, name, queue, test.expectedPriorityItems)
 	}
 }
 
@@ -177,93 +201,10 @@ func TestPriorityQueue_Len(t *testing.T) {
 			queue.Push(item, test.priorities[i])
 		}
 
-		len := queue.Len()
-
-		if len != test.expectedLen {
-			t.Errorf("%s: length expectancy doesn't match: expected %v, got %v", name, test.expectedLen, len)
-		}
-	}
-}
-
-func TestPriorityQueue_insertItemAt(t *testing.T) {
-	tests := map[string]struct {
-		items                 []priorityItem[int]
-		expectedPriorityItems []priorityItem[int]
-		insertItem            int
-		insertPriority        float64
-		insertIndex           int
-	}{
-		"insert in the middle of the queue": {
-			items: []priorityItem[int]{
-				{value: 10, priority: 10},
-				{value: 20, priority: 20},
-				{value: 30, priority: 30},
-				{value: 40, priority: 40},
-			},
-			insertItem:     25,
-			insertPriority: 25,
-			insertIndex:    2,
-			expectedPriorityItems: []priorityItem[int]{
-				{value: 10, priority: 10},
-				{value: 20, priority: 20},
-				{value: 25, priority: 25},
-				{value: 30, priority: 30},
-				{value: 40, priority: 40},
-			},
-		},
-		"insert at the start of the queue": {
-			items: []priorityItem[int]{
-				{value: 10, priority: 10},
-				{value: 20, priority: 20},
-				{value: 30, priority: 30},
-				{value: 40, priority: 40},
-			},
-			insertItem:     5,
-			insertPriority: 5,
-			insertIndex:    0,
-			expectedPriorityItems: []priorityItem[int]{
-				{value: 5, priority: 5},
-				{value: 10, priority: 10},
-				{value: 20, priority: 20},
-				{value: 30, priority: 30},
-				{value: 40, priority: 40},
-			},
-		},
-		"insert at the end of the queue": {
-			items: []priorityItem[int]{
-				{value: 10, priority: 10},
-				{value: 20, priority: 20},
-				{value: 30, priority: 30},
-				{value: 40, priority: 40},
-			},
-			insertItem:     50,
-			insertPriority: 50,
-			insertIndex:    4,
-			expectedPriorityItems: []priorityItem[int]{
-				{value: 10, priority: 10},
-				{value: 20, priority: 20},
-				{value: 30, priority: 30},
-				{value: 40, priority: 40},
-				{value: 50, priority: 50},
-			},
-		},
-	}
-
-	for name, test := range tests {
-		queue := &priorityQueue[int]{
-			items: test.items,
-		}
-
-		queue.insertItemAt(test.insertItem, test.insertPriority, test.insertIndex)
-
-		if len(queue.items) != len(test.expectedPriorityItems) {
-			t.Fatalf("%s: item length expectancy doesn't match: expected %v, got %v", name, len(test.expectedPriorityItems), len(queue.items))
-		}
+		n := queue.Len()
 
-		for i, expectedPriorityItem := range test.expectedPriorityItems {
-			if queue.items[i] != expectedPriorityItem {
-				t.Errorf("%s: item doesn't match: expected %v at index %d, got %v", name, expectedPriorityItem, i, queue.items[i])
-			}
+		if n != test.expectedLen {
+			t.Errorf("%s: length expectancy doesn't match: expected %v, got %v", name, test.expectedLen, n)
 		}
 	}
 }