switch to 0 based index

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
 	}
 }
 

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))
 		}