heapqueue/pkg/queue/queue.go

package queue

import (
	"fmt"
	"reflect"

	"git.wens.org.uk/heapqueue/pkg/heap"
	"golang.org/x/exp/constraints"
)

//QueueType is an enum for min/max queue
type QueueType bool

const (
	MINQUEUE QueueType = false
	MAXQUEUE QueueType = true
)

// Queue data type, holding a slice of generic queue objects 
type Queue[Orderable constraints.Ordered, T any] struct {
	store     []QueueObject[Orderable, T]
	queueSize int
}

// QueueObject holds generic types for key (comparable with > and < only), and any for value
type QueueObject[Orderable constraints.Ordered, T any] struct {
	Key   Orderable
	Value T
}

// BuildQueue takes an input slice of queue objects and the type of queue, returning a pointer to the constructed priority queue
func BuildQueue[Orderable constraints.Ordered, T any](d []QueueObject[Orderable, T], qType QueueType) *Queue[Orderable, T] {
	q := &Queue[Orderable, T]{d, len(d)}
	for i := len(d) / 2; i >= 0; i-- {
		if qType == MAXQUEUE {
			maxHeapify(q, i)
		} else {
			minHeapify(q, i)
		}
	}
	return q
}

// exchange two items within the context of the queue. requires max/min heapify afterwards
func exchange[Orderable constraints.Ordered, T any](q *Queue[Orderable, T], x, y int) {
	q.store[x], q.store[y] = q.store[y], q.store[x]
}

// maxHeapify reorders into a max heap starting at the input index i
func maxHeapify[Orderable constraints.Ordered, T any](q *Queue[Orderable, T], i int) {
	l, r := heap.Left(i), heap.Right(i)

	var largest int
	if l < q.queueSize && q.store[l].Key > q.store[i].Key {
		largest = l
	} else {
		largest = i
	}
	if r < q.queueSize && q.store[r].Key > q.store[largest].Key {
		largest = r
	}

	if largest != i {
		exchange(q, i, largest)
		maxHeapify(q, largest)
	}
}

// minHeapify reorders into a min heap starting at the input index i
func minHeapify[Orderable constraints.Ordered, T any](q *Queue[Orderable, T], i int) {
	l, r := heap.Left(i), heap.Right(i)

	var smallest int
	if l < q.queueSize && q.store[l].Key < q.store[i].Key {
		smallest = l
	} else {
		smallest = i
	}
	if r < q.queueSize && q.store[r].Key < q.store[smallest].Key {
		smallest = r
	}

	if smallest != i {
		exchange(q, i, smallest)
		minHeapify(q, smallest)
	}
}

// Maximum returns the QueueObject with the highest priority
func Maximum[Orderable constraints.Ordered, T any](q *Queue[Orderable, T]) (ret QueueObject[Orderable, T], e error) {
	if q.queueSize < 1 {
		e = fmt.Errorf("heap underflow")
		return
	}

	ret = q.store[0]
	return
}

// ExtractMaximum returns the QueueObject with the highest priority and removes it from the queue, reordering afterwards
func ExtractMaximum[Orderable constraints.Ordered, T any](q *Queue[Orderable, T]) (ret QueueObject[Orderable, T], e error) {
	ret, e = Maximum(q)
	if e != nil {
		return
	}

    q.store[0] = q.store[q.queueSize-1]
	q.queueSize--
    q.store = q.store[0:q.queueSize]

	maxHeapify(q, 0)

	return
}

// IncraseKey raises the priority of a given queue item and reorders the queue
func IncreaseKey[Orderable constraints.Ordered, T any](q *Queue[Orderable, T], item QueueObject[Orderable, T], newKey Orderable) error {
	if newKey < item.Key {
		return fmt.Errorf("new key %v is smaller than current key %v", newKey, item.Key)
	}

	var i int
	for idx, qObj := range q.store {
		if reflect.DeepEqual(qObj, item) {
			q.store[idx] = QueueObject[Orderable, T]{Key: newKey, Value: qObj.Value}
			i = idx
		}
	}

	for q.store[heap.Parent(i)].Key < q.store[i].Key {
		exchange(q, i, heap.Parent(i))
		i = heap.Parent(i)
	}

	return nil
}

// Insert adds a new item to the queue
func Insert[Orderable constraints.Ordered, T any](q *Queue[Orderable, T], item QueueObject[Orderable, T]) {
    q.queueSize++
    k := item.Key
    insert := QueueObject[Orderable, T]{Value: item.Value}
    q.store = append(q.store, insert)
    IncreaseKey(q, insert, k)
}

// IsMaxQueue returns true when the queue satisfies the max queue property i.e. holds a valid structure for a max queue
func IsMaxQueue[Orderable constraints.Ordered, T any](q *Queue[Orderable, T]) bool {
	for index, val := range q.store {
		if q.store[heap.Parent(index)].Key < val.Key {
			return false
		}
	}
	return true
}