golang 最小堆获取前 n 的数据

背景

大量数据，想获取其中 Num 降序前五的数据

实现

package testimport ("container/heap""fmt""testing"
)type Element struct {Content stringNum     int
}// 定义一个最小堆
type ElementMinHeap []Element// 重写方法
func (h ElementMinHeap) Len() int            { return len(h) }
func (h ElementMinHeap) Less(i, j int) bool  { return h[i].Num < h[j].Num }
func (h ElementMinHeap) Swap(i, j int)       { h[i], h[j] = h[j], h[i] }
func (h *ElementMinHeap) Push(x interface{}) { *h = append(*h, x.(Element)) }
func (h *ElementMinHeap) Pop() interface{} {old := *hn := len(old)x := old[n-1]*h = old[0 : n-1]return x
}func TestMinHeap(t *testing.T) {elements := []cronjob.Element{{"content1", 2},{"content2", 1},{"content3", 5},{"content4", 3},{"content5", 4},{"content6", 6},{"content7", 8},{"content8", 7},{"content9", 9},{"content10", 0},}topN := 5h := &cronjob.ElementMinHeap{}heap.Init(h)for _, elem := range elements {if h.Len() < topN {heap.Push(h, elem)} else if elem.Num > (*h)[0].Num {heap.Pop(h)heap.Push(h, elem)}}topElements := make([]cronjob.Element, h.Len())for i := h.Len() - 1; i >= 0; i-- {topElements[i] = heap.Pop(h).(cronjob.Element)}fmt.Println("Top", topN, "elements:")for _, elem := range topElements {fmt.Printf("Content: %s, Num: %d\n", elem.Content, elem.Num)}
}