github.com/bigcommerce/nomad@v0.9.3-bc/lib/kheap/score_heap.go (about)

     1  package kheap
     2  
     3  import (
     4  	"container/heap"
     5  )
     6  
     7  // HeapItem is an interface type implemented by objects stored in the ScoreHeap
     8  type HeapItem interface {
     9  	Data() interface{} // The data object
    10  	Score() float64    // Score to use as the sort criteria
    11  }
    12  
    13  // A ScoreHeap implements heap.Interface and is a min heap
    14  // that keeps the top K elements by Score. Push can be called
    15  // with an arbitrary number of values but only the top K are stored
    16  type ScoreHeap struct {
    17  	items    []HeapItem
    18  	capacity int
    19  }
    20  
    21  func NewScoreHeap(capacity uint32) *ScoreHeap {
    22  	return &ScoreHeap{capacity: int(capacity)}
    23  }
    24  
    25  func (pq ScoreHeap) Len() int { return len(pq.items) }
    26  
    27  func (pq ScoreHeap) Less(i, j int) bool {
    28  	return pq.items[i].Score() < pq.items[j].Score()
    29  }
    30  
    31  func (pq ScoreHeap) Swap(i, j int) {
    32  	pq.items[i], pq.items[j] = pq.items[j], pq.items[i]
    33  }
    34  
    35  // Push implements heap.Interface and only stores
    36  // the top K elements by Score
    37  func (pq *ScoreHeap) Push(x interface{}) {
    38  	item := x.(HeapItem)
    39  	if len(pq.items) < pq.capacity {
    40  		pq.items = append(pq.items, item)
    41  	} else {
    42  		// Pop the lowest scoring element if this item's Score is
    43  		// greater than the min Score so far
    44  		minIndex := 0
    45  		min := pq.items[minIndex]
    46  		if item.Score() > min.Score() {
    47  			// Replace min and heapify
    48  			pq.items[minIndex] = item
    49  			heap.Fix(pq, minIndex)
    50  		}
    51  	}
    52  }
    53  
    54  // Push implements heap.Interface and returns the top K scoring
    55  // elements in increasing order of Score. Callers must reverse the order
    56  // of returned elements to get the top K scoring elements in descending order
    57  func (pq *ScoreHeap) Pop() interface{} {
    58  	old := pq.items
    59  	n := len(old)
    60  	item := old[n-1]
    61  	pq.items = old[0 : n-1]
    62  	return item
    63  }
    64  
    65  // GetItemsReverse returns the items in this min heap in reverse order
    66  // sorted by score descending
    67  func (pq *ScoreHeap) GetItemsReverse() []interface{} {
    68  	ret := make([]interface{}, pq.Len())
    69  	i := pq.Len() - 1
    70  	for pq.Len() > 0 {
    71  		item := heap.Pop(pq)
    72  		ret[i] = item
    73  		i--
    74  	}
    75  	return ret
    76  }