mirror of
https://github.com/duke-git/lancet.git
synced 2026-02-04 12:52:28 +08:00
123 lines
2.8 KiB
Go
123 lines
2.8 KiB
Go
// Copyright 2021 dudaodong@gmail.com. All rights reserved.
|
|
// Use of this source code is governed by MIT license
|
|
|
|
// Package datastructure implements some data structure. eg. list, linklist, stack, queue, tree, graph.
|
|
package datastructure
|
|
|
|
import (
|
|
"github.com/duke-git/lancet/v2/lancetconstraints"
|
|
)
|
|
|
|
// MaxHeap implements a binary max heap
|
|
// type T should implements Compare function in lancetconstraints.Comparator interface.
|
|
type MaxHeap[T any] struct {
|
|
data []T
|
|
comparator lancetconstraints.Comparator
|
|
}
|
|
|
|
// NewMaxHeap returns a MaxHeap instance with the given comparator.
|
|
func NewMaxHeap[T any](comparator lancetconstraints.Comparator) *MaxHeap[T] {
|
|
return &MaxHeap[T]{
|
|
data: make([]T, 0),
|
|
comparator: comparator,
|
|
}
|
|
}
|
|
|
|
// Push value into the heap
|
|
func (h *MaxHeap[T]) Push(value T) {
|
|
h.data = append(h.data, value)
|
|
h.heapifyUp(len(h.data) - 1)
|
|
}
|
|
|
|
// heapifyUp heapify the data from bottom to top
|
|
func (h *MaxHeap[T]) heapifyUp(i int) {
|
|
for h.comparator.Compare(h.data[parentIndex(i)], h.data[i]) < 0 {
|
|
h.swap(parentIndex(i), i)
|
|
i = parentIndex(i)
|
|
}
|
|
}
|
|
|
|
// Pop return the largest value, and remove it from the heap
|
|
// if heap is empty, return zero value and fasle
|
|
func (h *MaxHeap[T]) Pop(value T) (T, bool) {
|
|
var val T
|
|
if h.Size() == 0 {
|
|
return val, false
|
|
}
|
|
|
|
val = h.data[0]
|
|
l := len(h.data) - 1
|
|
|
|
h.data[0] = h.data[l]
|
|
h.data = h.data[:l]
|
|
h.heapifyDown(0)
|
|
|
|
return val, true
|
|
}
|
|
|
|
// heapifyDown heapify the data from top to bottom
|
|
func (h *MaxHeap[T]) heapifyDown(i int) {
|
|
lastIndex := len(h.data) - 1
|
|
l, r := leftChildIndex(i), rightChildIndex(i)
|
|
childToCompare := 0
|
|
|
|
for l <= lastIndex {
|
|
if l == lastIndex {
|
|
childToCompare = l
|
|
} else if h.comparator.Compare(h.data[l], h.data[r]) > 0 {
|
|
childToCompare = l
|
|
} else {
|
|
childToCompare = r
|
|
}
|
|
|
|
if h.comparator.Compare(h.data[i], h.data[childToCompare]) < 0 {
|
|
h.swap(i, childToCompare)
|
|
i = childToCompare
|
|
l, r = leftChildIndex(i), rightChildIndex(i)
|
|
} else {
|
|
break
|
|
}
|
|
}
|
|
}
|
|
|
|
// Peek returns the largest element from the heap without removing it.
|
|
// if heap is empty, it returns zero value and false.
|
|
func (h *MaxHeap[T]) Peek() (T, bool) {
|
|
if h.Size() == 0 {
|
|
var val T
|
|
return val, false
|
|
}
|
|
|
|
return h.data[0], true
|
|
}
|
|
|
|
// Size return the number of elements in the heap
|
|
func (h *MaxHeap[T]) Size() int {
|
|
return len(h.data)
|
|
}
|
|
|
|
// Data return data of the heap
|
|
func (h *MaxHeap[T]) Data() []T {
|
|
return h.data
|
|
}
|
|
|
|
// parentIndex get parent index of the given index
|
|
func parentIndex(i int) int {
|
|
return (i - 1) / 2
|
|
}
|
|
|
|
// leftChildIndex get left child index of the given index
|
|
func leftChildIndex(i int) int {
|
|
return 2*i + 1
|
|
}
|
|
|
|
// rightChildIndex get right child index of the given index
|
|
func rightChildIndex(i int) int {
|
|
return 2*i + 2
|
|
}
|
|
|
|
// swap two elements in the heap
|
|
func (h *MaxHeap[T]) swap(i, j int) {
|
|
h.data[i], h.data[j] = h.data[j], h.data[i]
|
|
}
|