github.com/wfusion/gofusion@v1.1.14/common/utils/heap.go (about)

     1  package utils
     2  
     3  import (
     4  	"container/heap"
     5  	"errors"
     6  
     7  	"github.com/wfusion/gofusion/common/utils/clone"
     8  )
     9  
    10  var (
    11  	ErrOutOfRange = errors.New("out of range")
    12  )
    13  
    14  type _heap[E any] struct {
    15  	d   []E
    16  	cmp func(e1, e2 E) bool
    17  }
    18  
    19  func (h *_heap[E]) Len() int           { return len(h.d) }
    20  func (h *_heap[E]) Less(i, j int) bool { return h.cmp(h.d[i], h.d[j]) }
    21  func (h *_heap[E]) Swap(i, j int)      { h.d[i], h.d[j] = h.d[j], h.d[i] }
    22  func (h *_heap[E]) Push(x any)         { v := append(h.d, x.(E)); h.d = v }
    23  func (h *_heap[E]) Pop() (x any)       { x, h.d = h.d[len(h.d)-1], h.d[0:len(h.d)-1]; return x }
    24  
    25  // Heap base on generics to build a heap tree for any type
    26  type Heap[E any] struct {
    27  	data *_heap[E]
    28  }
    29  
    30  // Push pushes the element x onto the heap.
    31  // The complexity is O(log n) where n = h.Len().
    32  func (h *Heap[E]) Push(v E) { heap.Push(h.data, v) }
    33  
    34  // Pop removes and returns the minimum element (according to Less) from the heap.
    35  // The complexity is O(log n) where n = h.Len().
    36  // Pop is equivalent to Remove(h, 0).
    37  func (h *Heap[E]) Pop() E { return heap.Pop(h.data).(E) }
    38  
    39  func (h *Heap[E]) Element(index int) (e E, err error) {
    40  	if index < 0 || index >= h.data.Len() {
    41  		return e, ErrOutOfRange
    42  	}
    43  	return h.data.d[index], nil
    44  }
    45  
    46  // Remove removes and returns the element at index i from the heap.
    47  // The complexity is O(log n) where n = h.Len().
    48  func (h *Heap[E]) Remove(index int) E         { return heap.Remove(h.data, index).(E) }
    49  func (h *Heap[E]) Len() int                   { return len(h.data.d) }
    50  func (h *Heap[E]) Data() []E                  { return h.data.d }
    51  func (h *Heap[E]) Fix(index int)              { heap.Fix(h.data, index) }
    52  func (h *Heap[E]) Init()                      { heap.Init(h.data) }
    53  func (h *Heap[E]) CmpFn() func(e1, e2 E) bool { return h.data.cmp }
    54  
    55  // Clone heap
    56  func (h *Heap[E]) Clone() (r *Heap[E]) {
    57  	if h == nil {
    58  		return
    59  	}
    60  	ret := &_heap[E]{cmp: h.data.cmp, d: make([]E, 0, len(h.data.d))}
    61  	for _, e := range h.data.d {
    62  		if elem, ok := any(e).(clonable[E]); ok {
    63  			ret.d = append(ret.d, elem.Clone())
    64  		} else {
    65  			ret.d = append(ret.d, clone.Clone(e))
    66  		}
    67  	}
    68  	heap.Init(ret)
    69  	return &Heap[E]{data: ret}
    70  }
    71  
    72  // Copy heap
    73  func (h *Heap[E]) Copy() (r *Heap[E]) {
    74  	if h == nil {
    75  		return
    76  	}
    77  	ret := &_heap[E]{cmp: h.data.cmp}
    78  	ret.d = make([]E, len(h.data.d))
    79  	copy(ret.d, h.data.d)
    80  	heap.Init(ret)
    81  	return &Heap[E]{data: ret}
    82  }
    83  
    84  // NewHeap return an initial Heap pointer
    85  func NewHeap[E any](t []E, cmp func(e1, e2 E) bool) *Heap[E] {
    86  	ret := &_heap[E]{d: t, cmp: cmp}
    87  	heap.Init(ret)
    88  	return &Heap[E]{data: ret}
    89  }