github.com/lingyao2333/mo-zero@v1.4.1/core/stat/topk.go (about)

     1  package stat
     2  
     3  import "container/heap"
     4  
     5  type taskHeap []Task
     6  
     7  func (h *taskHeap) Len() int {
     8  	return len(*h)
     9  }
    10  
    11  func (h *taskHeap) Less(i, j int) bool {
    12  	return (*h)[i].Duration < (*h)[j].Duration
    13  }
    14  
    15  func (h *taskHeap) Swap(i, j int) {
    16  	(*h)[i], (*h)[j] = (*h)[j], (*h)[i]
    17  }
    18  
    19  func (h *taskHeap) Push(x interface{}) {
    20  	*h = append(*h, x.(Task))
    21  }
    22  
    23  func (h *taskHeap) Pop() interface{} {
    24  	old := *h
    25  	n := len(old)
    26  	x := old[n-1]
    27  	*h = old[0 : n-1]
    28  	return x
    29  }
    30  
    31  func topK(all []Task, k int) []Task {
    32  	h := new(taskHeap)
    33  	heap.Init(h)
    34  
    35  	for _, each := range all {
    36  		if h.Len() < k {
    37  			heap.Push(h, each)
    38  		} else if (*h)[0].Duration < each.Duration {
    39  			heap.Pop(h)
    40  			heap.Push(h, each)
    41  		}
    42  	}
    43  
    44  	return *h
    45  }