You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 

59 lines
1.0 KiB

package container
import (
"cmp"
"container/heap"
)
// 堆封装
type heapQueue[T any, C cmp.Ordered] struct {
queue []T
compare func(T) C
}
func (h *heapQueue[T, C]) Less(i, j int) bool {
return h.compare(h.queue[i]) < h.compare(h.queue[j])
}
func (h *heapQueue[T, C]) Swap(i, j int) {
h.queue[i], h.queue[j] = h.queue[j], h.queue[i]
}
func (h *heapQueue[T, C]) Len() int {
return len(h.queue)
}
func (h *heapQueue[T, C]) Pop() (v any) {
h.queue, v = h.queue[:h.Len()-1], h.queue[h.Len()-1]
return
}
func (h *heapQueue[T, C]) Push(v any) {
h.queue = append(h.queue, v.(T))
}
// Heap 封装堆操作
type Heap[T any, C cmp.Ordered] struct {
queue *heapQueue[T, C]
}
func NewHeap[T any, C cmp.Ordered](compare func(T) C) *Heap[T, C] {
return &Heap[T, C]{
queue: &heapQueue[T, C]{compare: compare},
}
}
func (h *Heap[T, C]) Push(v T) {
heap.Push(h.queue, v)
}
func (h *Heap[T, C]) Pop() T {
return heap.Pop(h.queue).(T)
}
func (h *Heap[T, C]) Peek() (v T, ok bool) {
if h.queue.Len() == 0 {
return v, false
}
return h.queue.queue[0], true
}