golang 最小堆获取前 n 的数据

背景

大量数据，想获取其中 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)
	}
}