diff --git a/sorting/priority_queue/index_priority_queue.go b/sorting/priority_queue/index_priority_queue.go index 6e59845..ba6853d 100644 --- a/sorting/priority_queue/index_priority_queue.go +++ b/sorting/priority_queue/index_priority_queue.go @@ -1,6 +1,5 @@ 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 @@ -23,7 +22,9 @@ type indexPriorityQueue[T any] struct { pq []int // "reverse" for pq. Maps item key to priority // qp[key] = priority index, qp[pq[key]] = pq[qp[key]] = key - qp []int + qp []int + // "less" function. Depending on desired order first element + // not necessary less than a second less func(T, T) bool } @@ -31,10 +32,9 @@ type indexPriorityQueue[T any] struct { // 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) + items := make([]T, indexSize) + pq := make([]int, indexSize) + qp := make([]int, indexSize) for i := range qp { qp[i] = -1 @@ -48,42 +48,35 @@ func NewIPQ[T any](less func(T, T) bool, indexSize int) IndexPriorityQueue[T] { } func (q *indexPriorityQueue[T]) top() T { - return q.items[q.pq[1]] + return q.items[q.pq[0]] } func (q *indexPriorityQueue[T]) topKey() int { - return q.pq[1] + return q.pq[0] } 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) + 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] + return q.removeKey(q.topKey()) } func (q *indexPriorityQueue[T]) removeKey(key int) T { pivot := q.qp[key] - q.swap(pivot, q.n) q.n-- + q.swap(pivot, 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 + q.pq[q.n] = -1 return q.items[key] } @@ -106,23 +99,40 @@ func (q *indexPriorityQueue[_]) contains(key int) bool { } 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]]) { + + for { + child := 2*parent + 1 + + if child >= q.n { + break + } + + if child+1 < 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 + for { + parent := (child - 1) / 2 + + if child <= 0 || q.less(q.items[q.pq[child]], q.items[q.pq[parent]]) { + break + } + + q.swap(parent, child) + + child = parent } } diff --git a/sorting/priority_queue/index_priority_queue_test.go b/sorting/priority_queue/index_priority_queue_test.go index d0759ae..3220b14 100644 --- a/sorting/priority_queue/index_priority_queue_test.go +++ b/sorting/priority_queue/index_priority_queue_test.go @@ -170,18 +170,18 @@ func TestMultiwayMerge(t *testing.T) { expected := "AABBBCDEFFGHIIJNPQQZ" actual := "" - pq := NewIPQ(func(t1, t2 string) bool { return t1 > t2 }, len(allStreams)+1) + pq := NewIPQ(func(t1, t2 string) bool { return t1 > t2 }, len(allStreams)) for i, stream := range allStreams { rune, _, _ := stream.ReadRune() - pq.insert(i+1, string(rune)) + pq.insert(i, string(rune)) } for !pq.isEmpty() { actual += string(pq.top()) streamIndex := pq.topKey() pq.remove() - rune, _, err := allStreams[streamIndex-1].ReadRune() + rune, _, err := allStreams[streamIndex].ReadRune() if err != io.EOF { pq.insert(streamIndex, string(rune)) }