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.

128 lines
3.1 KiB
Go

package heap
import "golang.org/x/exp/constraints"
type HeapType bool
const (
MINHEAP HeapType = false
MAXHEAP HeapType = true
)
// Heap object for generic orderable (i.e. can be compared with < or >) data types
type Heap[Orderable constraints.Ordered] struct {
data []Orderable
heapSize int
}
// Parent index is returned for an input index
func Parent(index int) (parent int) {
return (index - 1) / 2
}
// Left index is returned for an input index
func Left(index int) (left int) {
return 2*index + 1
}
// Right index is returned for an input index
func Right(index int) (right int) {
return 2*index + 2
}
// exchange two items within the context of the heap. requires max/min heapify afterwards
func exchange[Orderable constraints.Ordered](h *Heap[Orderable], x, y int) {
h.data[x], h.data[y] = h.data[y], h.data[x]
}
// IsMaxHeap returns true when the heap satisfies the max heap property i.e. holds a valid structure for a max heap
func IsMaxHeap[Orderable constraints.Ordered](h *Heap[Orderable]) bool {
for index, val := range h.data {
if h.data[Parent(index)] < val {
return false
}
}
return true
}
// IsMinHeap returns true when the heap satisfies the min heap property i.e. holds a valid structure for a min heap
func IsMinHeap[Orderable constraints.Ordered](h *Heap[Orderable]) bool {
for index, val := range h.data {
if h.data[Parent(index)] > val {
return false
}
}
return true
}
// maxHeapify reorders into a min heap starting at the input index i
func maxHeapify[Orderable constraints.Ordered](h *Heap[Orderable], i int) {
l, r := Left(i), Right(i)
var largest int
if l < h.heapSize && h.data[l] > h.data[i] {
largest = l
} else {
largest = i
}
if r < h.heapSize && h.data[r] > h.data[largest] {
largest = r
}
if largest != i {
exchange(h, i, largest)
maxHeapify(h, largest)
}
}
// minHeapify reorders into a min heap starting at the input index i
func minHeapify[Orderable constraints.Ordered](h *Heap[Orderable], i int) {
l, r := Left(i), Right(i)
var smallest int
if l < h.heapSize && h.data[l] < h.data[i] {
smallest = l
} else {
smallest = i
}
if r < h.heapSize && h.data[r] < h.data[smallest] {
smallest = r
}
if smallest != i {
exchange(h, i, smallest)
minHeapify(h, smallest)
}
}
// BuildHeap takes a generic slice (any type that can be compared with < and >) and the type of heap and returns a pointer to a heap structure
func BuildHeap[Orderable constraints.Ordered](d []Orderable, maxHeap HeapType) *Heap[Orderable] {
heap := &Heap[Orderable]{d, len(d)}
for i := len(d) / 2; i >= 0; i-- {
if maxHeap {
maxHeapify(heap, i)
} else {
minHeapify(heap, i)
}
}
return heap
}
// HeapSort takes a generic slice (any type that can be compared with < and >) and returns the sorted slice
func HeapSort[Orderable constraints.Ordered](d []Orderable, maxHeap HeapType) []Orderable {
heap := BuildHeap(d, maxHeap)
for i := len(d) - 1; i > 0; i-- {
exchange(heap, 0, i)
heap.heapSize = heap.heapSize - 1
if maxHeap {
maxHeapify(heap, 0)
} else {
minHeapify(heap, 0)
}
}
return heap.data
}