indexed priority queue

sorting/priority_queue/index_priority_queue.go

@@ -0,0 +1,133 @@
+package priority_queue
+
+// TODO: change name to "keyed priority queue"
+type IndexPriorityQueue[T any] interface {
+	top() T                 // get item with biggest priority
+	topKey() int            // get key of an item with biggest priority
+	remove() T              // removes item with the biggest priority
+	removeKey(key int) T    // removes item with specified key
+	insert(key int, item T) // adds item with specified key
+	change(key int, item T) // changes item with specified key
+	contains(key int) bool  // checks if key exists in queue
+	isEmpty() bool
+	size() int
+}
+
+type indexPriorityQueue[T any] struct {
+	n int
+	// unordered items.
+	// items[key] = item
+	items []T
+	// priority queue. Contains keys for items in priority order.
+	// items[pq[1]] = item with biggest priority
+	pq []int
+	// "reverse" for pq. Maps item key to priority
+	// qp[key] = priority index, qp[pq[key]] = pq[qp[key]] = key
+	qp   []int
+	less func(T, T) bool
+}
+
+// TODO: panic for illegal operations
+// TODO: can we construct queue without bounded index size?
+func NewIPQ[T any](less func(T, T) bool, indexSize int) IndexPriorityQueue[T] {
+	n := 0
+	// TODO: switch to 0 based index
+	items := make([]T, indexSize+1)
+	pq := make([]int, indexSize+1)
+	qp := make([]int, indexSize+1)
+
+	for i := range qp {
+		qp[i] = -1
+	}
+
+	for i := range pq {
+		pq[i] = -1
+	}
+
+	return &indexPriorityQueue[T]{n, items, pq, qp, less}
+}
+
+func (q *indexPriorityQueue[T]) top() T {
+	return q.items[q.pq[1]]
+}
+
+func (q *indexPriorityQueue[T]) topKey() int {
+	return q.pq[1]
+}
+
+func (q *indexPriorityQueue[T]) insert(key int, item T) {
+	q.n++
+	q.pq[q.n] = key
+	q.qp[key] = q.n
+	q.items[key] = item
+	q.swim(q.n)
+}
+
+func (q *indexPriorityQueue[T]) remove() T {
+	topKey := q.topKey()
+	q.swap(1, q.n)
+	q.n--
+	q.sink(1)
+	// TODO: need to remove actual item from items array
+	// to allow removed item to be GCed
+	q.qp[q.pq[q.n+1]] = -1
+	return q.items[topKey]
+}
+
+func (q *indexPriorityQueue[T]) removeKey(key int) T {
+	pivot := q.qp[key]
+	q.swap(pivot, q.n)
+	q.n--
+	q.swim(pivot)
+	q.sink(pivot)
+	// TODO: need to remove actual item from items array
+	// to prevent memory leak
+	q.qp[key] = -1
+	q.pq[q.n+1] = -1
+	return q.items[key]
+}
+
+func (q *indexPriorityQueue[T]) change(key int, item T) {
+	q.items[key] = item
+	q.swim(q.qp[key])
+	q.sink(q.qp[key])
+}
+
+func (pq *indexPriorityQueue[_]) size() int {
+	return pq.n
+}
+
+func (pq *indexPriorityQueue[_]) isEmpty() bool {
+	return pq.n == 0
+}
+
+func (q *indexPriorityQueue[_]) contains(key int) bool {
+	return q.qp[key] != -1
+}
+
+func (q *indexPriorityQueue[T]) sink(parent int) {
+	for 2*parent <= q.n {
+		child := 2 * parent
+		if child < q.n && q.less(q.items[q.pq[child]], q.items[q.pq[child+1]]) {
+			child++
+		}
+		if !q.less(q.items[q.pq[parent]], q.items[q.pq[child]]) {
+			break
+		}
+		q.swap(parent, child)
+		parent = child
+	}
+}
+
+func (q *indexPriorityQueue[T]) swim(child int) {
+	for child > 1 && q.less(q.items[q.pq[child/2]], q.items[q.pq[child]]) {
+		q.swap(child/2, child)
+		child = child / 2
+	}
+}
+
+func (q *indexPriorityQueue[_]) swap(i, j int) {
+	q.qp[q.pq[i]] = j
+	q.qp[q.pq[j]] = i
+	q.pq[i], q.pq[j] = q.pq[j], q.pq[i]
+}

sorting/priority_queue/index_priority_queue_test.go

@@ -0,0 +1,158 @@
+package priority_queue
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNewIPQ(t *testing.T) {
+	q := NewIPQ(intDescending, 1)
+	assert.NotNil(t, q)
+
+	assert.Equal(t, 0, q.size())
+	assert.Equal(t, true, q.isEmpty())
+	assert.Equal(t, -1, q.topKey())
+	// TODO: maybe should return nil?
+	// assert.Equal(t, 0, q.top())
+}
+
+func TestIPQInsert(t *testing.T) {
+	pq := NewIPQ(intDescending, 10)
+
+	pq.insert(1, 3)
+	assert.Equal(t, 3, pq.top())
+
+	pq.insert(2, 4)
+	assert.Equal(t, 4, pq.top())
+
+	pq.insert(3, 1)
+	assert.Equal(t, 4, pq.top())
+
+	pq.insert(4, 4)
+	assert.Equal(t, 4, pq.top())
+}
+
+func TestMoreIPQInsert(t *testing.T) {
+	pq := NewIPQ(intDescending, 10)
+
+	pq.insert(1, 10)
+
+	assert.Equal(t, 1, pq.topKey())
+	assert.Equal(t, 10, pq.top())
+	assert.Equal(t, 1, pq.size())
+	assert.Equal(t, true, pq.contains(1))
+	assert.Equal(t, false, pq.isEmpty())
+
+	pq.insert(2, 20)
+
+	assert.Equal(t, 2, pq.topKey())
+	assert.Equal(t, 20, pq.top())
+	assert.Equal(t, 2, pq.size())
+	assert.Equal(t, true, pq.contains(2))
+	assert.Equal(t, false, pq.isEmpty())
+}
+
+func TestIPQRemove(t *testing.T) {
+	pq := NewIPQ(intDescending, 10)
+
+	pq.insert(1, 10)
+
+	assert.Equal(t, 1, pq.topKey())
+	assert.Equal(t, 10, pq.top())
+	assert.Equal(t, 1, pq.size())
+	assert.Equal(t, true, pq.contains(1))
+	assert.Equal(t, false, pq.isEmpty())
+
+	pq.insert(2, 20)
+
+	assert.Equal(t, 2, pq.topKey())
+	assert.Equal(t, 20, pq.top())
+	assert.Equal(t, 2, pq.size())
+	assert.Equal(t, true, pq.contains(2))
+	assert.Equal(t, false, pq.isEmpty())
+
+	removed := pq.remove()
+
+	assert.Equal(t, 20, removed)
+	assert.Equal(t, 10, pq.top())
+	assert.Equal(t, 1, pq.size())
+	assert.Equal(t, false, pq.contains(2))
+	assert.Equal(t, true, pq.contains(1))
+	assert.Equal(t, false, pq.isEmpty())
+
+	removed = pq.remove()
+
+	assert.Equal(t, 10, removed)
+	// TODO: should return nil?
+	// assert.Equal(t, "", pq.top())
+	assert.Equal(t, 0, pq.size())
+	assert.Equal(t, false, pq.contains(1))
+	assert.Equal(t, true, pq.isEmpty())
+}
+
+func TestIPQRemoveAtIndex(t *testing.T) {
+	pq := NewIPQ(intDescending, 10)
+
+	// top -> 40 - 30 - 20 - 10
+	pq.insert(8, 10)
+	pq.insert(5, 20)
+	pq.insert(3, 30)
+	pq.insert(4, 40)
+
+	assert.Equal(t, 40, pq.top())
+	assert.Equal(t, 4, pq.topKey())
+
+	// top -> 40 - 30 - 10
+	removed := pq.removeKey(5)
+
+	assert.Equal(t, 20, removed)
+	assert.Equal(t, 40, pq.top())
+	assert.Equal(t, 4, pq.topKey())
+
+	// top -> 30 - 10
+	removed = pq.removeKey(4)
+
+	assert.Equal(t, 40, removed)
+	assert.Equal(t, 30, pq.top())
+	assert.Equal(t, 3, pq.topKey())
+
+	// top -> 30 - 20 - 10
+	pq.insert(5, 20)
+
+	assert.Equal(t, 30, pq.top())
+	assert.Equal(t, 3, pq.topKey())
+
+	// top -> 10
+	removed = pq.removeKey(3)
+	assert.Equal(t, 30, removed)
+	removed = pq.removeKey(5)
+	assert.Equal(t, 20, removed)
+
+	assert.Equal(t, 10, pq.top())
+	assert.Equal(t, 8, pq.topKey())
+	assert.Equal(t, 1, pq.size())
+	assert.Equal(t, false, pq.contains(5))
+	assert.Equal(t, false, pq.contains(4))
+	assert.Equal(t, false, pq.contains(3))
+	assert.Equal(t, true, pq.contains(8))
+	assert.Equal(t, false, pq.isEmpty())
+}
+
+func TestIndexChange(t *testing.T) {
+	pq := NewIPQ(intDescending, 10)
+
+	pq.insert(1, 9)
+	pq.insert(2, 8)
+	pq.insert(3, 12)
+
+	assert.Equal(t, 12, pq.top())
+
+	pq.change(3, 7)
+
+	assert.Equal(t, 9, pq.top())
+
+	pq.change(2, 10)
+
+	assert.Equal(t, 10, pq.top())
+}

sorting/priority_queue/priority_queue.go

@@ -1,4 +1,4 @@
-package sorting
+package priority_queue
 
 type PriorityQueue[T any] interface {
 	top() T
@@ -66,24 +66,6 @@ func (pq *priorityQueue[_]) isEmpty() bool {
 	return pq.n == 0
 }
 
-func (pq *priorityQueue[_]) swap(i, j int) {
-	pq.heap[i], pq.heap[j] = pq.heap[j], pq.heap[i]
-}
-
-func (pq *priorityQueue[T]) swim(child int) {
-	// Until we reach top of the heap
-	// and parent node is less than current child
-	for child > 1 && pq.less(pq.heap[child/2], pq.heap[child]) {
-
-		// We swap parent with the child
-		pq.swap(child/2, child)
-
-		// Parent node becomes new child
-		// for next iteration
-		child = child / 2
-	}
-}
-
 func (pq *priorityQueue[T]) sink(parent int) {
 	// While parent has some children
 	for 2*parent <= pq.n {
@@ -111,3 +93,21 @@ func (pq *priorityQueue[T]) sink(parent int) {
 		parent = child
 	}
 }
+
+func (pq *priorityQueue[T]) swim(child int) {
+	// Until we reach top of the heap
+	// and parent node is less than current child
+	for child > 1 && pq.less(pq.heap[child/2], pq.heap[child]) {
+
+		// We swap parent with the child
+		pq.swap(child/2, child)
+
+		// Parent node becomes new child
+		// for next iteration
+		child = child / 2
+	}
+}
+
+func (pq *priorityQueue[T]) swap(i, j int) {
+	pq.heap[i], pq.heap[j] = pq.heap[j], pq.heap[i]
+}

sorting/priority_queue/priority_queue_test.go

@@ -1,4 +1,4 @@
-package sorting
+package priority_queue
 
 import (
 	"testing"
@@ -6,15 +6,13 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-var intCompare = func(t1, t2 int) bool { return t1 < t2 }
-
 func TestNew(t *testing.T) {
-	pq := NewPQ(intCompare)
+	pq := NewPQ(intDescending)
 	assert.NotNil(t, pq)
 }
 
 func TestSize(t *testing.T) {
-	pq := NewPQ(intCompare)
+	pq := NewPQ(intDescending)
 	assert.Equal(t, 0, pq.size())
 	pq.insert(1)
 	assert.Equal(t, 1, pq.size())
@@ -24,14 +22,14 @@ func TestSize(t *testing.T) {
 }
 
 func TestIsEmpty(t *testing.T) {
-	pq := NewPQ(intCompare)
+	pq := NewPQ(intDescending)
 	assert.Equal(t, true, pq.isEmpty())
 	pq.insert(1)
 	assert.Equal(t, false, pq.isEmpty())
 }
 
 func TestInsert(t *testing.T) {
-	pq := NewPQ(intCompare)
+	pq := NewPQ(intDescending)
 	pq.insert(1)
 	assert.Equal(t, 1, pq.top())
 	pq.insert(4)
@@ -43,7 +41,7 @@ func TestInsert(t *testing.T) {
 }
 
 func TestDelete(t *testing.T) {
-	pq := NewPQ(intCompare)
+	pq := NewPQ(intDescending)
 	pq.insert(1)
 	pq.insert(2)
 	pq.insert(6)

sorting/priority_queue/testing.go

@@ -0,0 +1,3 @@
+package priority_queue
+
+var intDescending = func(t1, t2 int) bool { return t1 < t2 }