125 lines
2.3 KiB
Go
125 lines
2.3 KiB
Go
// Copyright (C) 2025 wangyusong
|
|
//
|
|
// This program is free software: you can redistribute it and/or modify
|
|
// it under the terms of the GNU Affero General Public License as published by
|
|
// the Free Software Foundation, either version 3 of the License, or
|
|
// (at your option) any later version.
|
|
//
|
|
// This program is distributed in the hope that it will be useful,
|
|
// but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
// GNU Affero General Public License for more details.
|
|
//
|
|
// You should have received a copy of the GNU Affero General Public License
|
|
// along with this program. If not, see <https://www.gnu.org/licenses/>.
|
|
|
|
package heap
|
|
|
|
import (
|
|
"container/heap"
|
|
"sort"
|
|
)
|
|
|
|
type Heap[T any] struct {
|
|
inner *innerHeap[T]
|
|
limit int
|
|
}
|
|
|
|
func New[T any](data []T, less func(a, b T) bool) *Heap[T] {
|
|
h := &Heap[T]{
|
|
inner: newInnerHeap(data, less),
|
|
limit: cap(data),
|
|
}
|
|
heap.Init(h.inner)
|
|
|
|
return h
|
|
}
|
|
|
|
func (h *Heap[T]) TryEvictPush(x T) {
|
|
switch {
|
|
case h.Len() < h.limit:
|
|
case h.inner.less(h.Peek(), x):
|
|
h.Pop()
|
|
default:
|
|
return
|
|
}
|
|
|
|
h.Push(x)
|
|
}
|
|
|
|
func (h *Heap[T]) Push(x T) {
|
|
heap.Push(h.inner, x)
|
|
}
|
|
|
|
func (h *Heap[T]) Pop() T {
|
|
return heap.Pop(h.inner).(T)
|
|
}
|
|
|
|
func (h *Heap[T]) PopLast() T {
|
|
return heap.Remove(h.inner, h.Len()-1).(T)
|
|
}
|
|
|
|
func (h *Heap[T]) Peek() T {
|
|
if h.Len() == 0 {
|
|
var zero T
|
|
|
|
return zero
|
|
}
|
|
|
|
return h.inner.data[0]
|
|
}
|
|
|
|
func (h *Heap[T]) Len() int {
|
|
return h.inner.Len()
|
|
}
|
|
|
|
func (h *Heap[T]) Cap() int {
|
|
return h.limit
|
|
}
|
|
|
|
func (h *Heap[T]) Slice() []T {
|
|
return h.inner.data
|
|
}
|
|
|
|
func (h *Heap[T]) DESCSort() {
|
|
sort.Slice(h.inner.data, func(i, j int) bool {
|
|
return !h.inner.less(h.inner.data[i], h.inner.data[j])
|
|
})
|
|
}
|
|
|
|
type innerHeap[T any] struct {
|
|
data []T
|
|
less func(a, b T) bool
|
|
}
|
|
|
|
func newInnerHeap[T any](data []T, less func(a, b T) bool) *innerHeap[T] {
|
|
return &innerHeap[T]{
|
|
data: data,
|
|
less: less,
|
|
}
|
|
}
|
|
|
|
func (h *innerHeap[T]) Len() int {
|
|
return len(h.data)
|
|
}
|
|
|
|
func (h *innerHeap[T]) Less(i, j int) bool {
|
|
return h.less(h.data[i], h.data[j])
|
|
}
|
|
|
|
func (h *innerHeap[T]) Swap(i, j int) {
|
|
h.data[i], h.data[j] = h.data[j], h.data[i]
|
|
}
|
|
|
|
func (h *innerHeap[T]) Push(x any) {
|
|
h.data = append(h.data, x.(T))
|
|
}
|
|
|
|
func (h *innerHeap[T]) Pop() any {
|
|
n := len(h.data)
|
|
x := h.data[n-1]
|
|
h.data = h.data[:n-1]
|
|
|
|
return x
|
|
}
|