diff --git a/sorting/priority_queue/index_priority_queue.go b/sorting/priority_queue/index_priority_queue.go new file mode 100644 index 0000000..6e59845 --- /dev/null +++ b/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] +} diff --git a/sorting/priority_queue/index_priority_queue_test.go b/sorting/priority_queue/index_priority_queue_test.go new file mode 100644 index 0000000..f98faaf --- /dev/null +++ b/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()) +} diff --git a/sorting/priority_queue.go b/sorting/priority_queue/priority_queue.go similarity index 97% rename from sorting/priority_queue.go rename to sorting/priority_queue/priority_queue.go index fea5b3c..fa44a02 100644 --- a/sorting/priority_queue.go +++ b/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] +} diff --git a/sorting/priority_queue_test.go b/sorting/priority_queue/priority_queue_test.go similarity index 82% rename from sorting/priority_queue_test.go rename to sorting/priority_queue/priority_queue_test.go index ee7d43a..540255a 100644 --- a/sorting/priority_queue_test.go +++ b/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) diff --git a/sorting/priority_queue/testing.go b/sorting/priority_queue/testing.go new file mode 100644 index 0000000..2f1d148 --- /dev/null +++ b/sorting/priority_queue/testing.go @@ -0,0 +1,3 @@ +package priority_queue + +var intDescending = func(t1, t2 int) bool { return t1 < t2 }