github.com/awesome-flow/flow@v0.0.3-0.20190918184116-508d75d68a2c/pkg/util/data/bin_heap.go (about)

     1  package data
     2  
     3  import (
     4  	"sync"
     5  )
     6  
     7  type BinHeap struct {
     8  	mutex *sync.Mutex
     9  	vals  []*BinHeapNode
    10  }
    11  
    12  type BinHeapNode struct {
    13  	weight uint32
    14  	value  interface{}
    15  }
    16  
    17  func NewBinHeap() *BinHeap {
    18  	return &BinHeap{
    19  		mutex: &sync.Mutex{},
    20  		vals:  make([]*BinHeapNode, 0),
    21  	}
    22  }
    23  
    24  func (hp *BinHeap) Insert(weight uint32, val interface{}) {
    25  	hp.mutex.Lock()
    26  	defer hp.mutex.Unlock()
    27  	hp.vals = append(hp.vals, &BinHeapNode{weight: weight, value: val})
    28  	hp.heapify()
    29  }
    30  
    31  func (hp *BinHeap) heapify() {
    32  	l := len(hp.vals) - 1
    33  	if l <= 0 {
    34  		return
    35  	}
    36  	for {
    37  		p := (l - 1) / 2
    38  		if hp.vals[l].weight >= hp.vals[p].weight {
    39  			// swap them
    40  			hp.vals[p], hp.vals[l] = hp.vals[l], hp.vals[p]
    41  		} else {
    42  			break
    43  		}
    44  		if p <= 0 {
    45  			break
    46  		}
    47  	}
    48  }
    49  
    50  func (hp *BinHeap) GetMax() interface{} {
    51  	hp.mutex.Lock()
    52  	defer hp.mutex.Unlock()
    53  	if len(hp.vals) > 0 {
    54  		return hp.vals[0].value
    55  	}
    56  	return nil
    57  }