You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
104 lines
2.6 KiB
Go
104 lines
2.6 KiB
Go
package container
|
|
|
|
import (
|
|
"golang.org/x/exp/constraints"
|
|
)
|
|
|
|
type HeapItem[T any, V constraints.Integer] struct {
|
|
Value T
|
|
Priority V
|
|
}
|
|
|
|
// Heap represents a generic heap data structure.
|
|
type Heap[T any, V constraints.Integer] struct {
|
|
elements []HeapItem[T, V]
|
|
isMinHeap bool
|
|
}
|
|
|
|
// NewHeap creates a new heap with the specified ordering.
|
|
// The isMinHeap parameter determines whether the heap is a min heap or a max heap.
|
|
// The heap is initially empty.
|
|
func NewHeap[T any, V constraints.Integer](isMinHeap bool) *Heap[T, V] {
|
|
return &Heap[T, V]{elements: []HeapItem[T, V]{}, isMinHeap: isMinHeap}
|
|
}
|
|
|
|
// Push adds an element to the heap.
|
|
func (h *Heap[T, V]) Push(element T, priority V) {
|
|
h.elements = append(h.elements, HeapItem[T, V]{element, priority})
|
|
h.heapifyUp()
|
|
}
|
|
|
|
// Pop removes and returns the root element from the heap.
|
|
func (h *Heap[T, V]) Pop() T {
|
|
if len(h.elements) == 0 {
|
|
var zero T // Return zero value of T
|
|
return zero
|
|
}
|
|
|
|
root := h.elements[0]
|
|
lastIndex := len(h.elements) - 1
|
|
h.elements[0] = h.elements[lastIndex]
|
|
h.elements = h.elements[:lastIndex]
|
|
|
|
h.heapifyDown()
|
|
return root.Value
|
|
}
|
|
|
|
// Peek returns the root element without removing it.
|
|
func (h *Heap[T, V]) Peek() T {
|
|
if len(h.elements) == 0 {
|
|
var zero T // Return zero value of T
|
|
return zero
|
|
}
|
|
return h.elements[0].Value
|
|
}
|
|
|
|
// heapifyUp adjusts the heap after adding a new element.
|
|
func (h *Heap[T, V]) heapifyUp() {
|
|
index := len(h.elements) - 1
|
|
for index > 0 {
|
|
parentIndex := (index - 1) / 2
|
|
if h.compare(h.elements[index].Priority, h.elements[parentIndex].Priority) {
|
|
h.elements[index], h.elements[parentIndex] = h.elements[parentIndex], h.elements[index]
|
|
index = parentIndex
|
|
} else {
|
|
break
|
|
}
|
|
}
|
|
}
|
|
|
|
// heapifyDown adjusts the heap after removing the root element.
|
|
func (h *Heap[T, V]) heapifyDown() {
|
|
index := 0
|
|
lastIndex := len(h.elements) - 1
|
|
for index < lastIndex {
|
|
leftChildIndex := 2*index + 1
|
|
rightChildIndex := 2*index + 2
|
|
|
|
var childIndex int
|
|
if leftChildIndex <= lastIndex {
|
|
childIndex = leftChildIndex
|
|
if rightChildIndex <= lastIndex && h.compare(h.elements[rightChildIndex].Priority, h.elements[leftChildIndex].Priority) {
|
|
childIndex = rightChildIndex
|
|
}
|
|
|
|
if h.compare(h.elements[childIndex].Priority, h.elements[index].Priority) {
|
|
h.elements[index], h.elements[childIndex] = h.elements[childIndex], h.elements[index]
|
|
index = childIndex
|
|
} else {
|
|
break
|
|
}
|
|
} else {
|
|
break
|
|
}
|
|
}
|
|
}
|
|
|
|
// compare compares two elements based on the heap type.
|
|
func (h *Heap[T, V]) compare(x, y V) bool {
|
|
if h.isMinHeap {
|
|
return x < y
|
|
}
|
|
return x > y
|
|
}
|