背景
大量数据,想获取其中 Num 降序前五的数据
实现
package test
import (
"container/heap"
"fmt"
"testing"
)
type Element struct {
Content string
Num 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 := *h
n := 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 := 5
h := &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)
}
}