限制Go堆接口实现的优先级队列的大小

如果您只需要N个元素中的前M个最小元素,那么请使用将为您提供最大元素的堆,并在其大小超过值M时修剪堆。然后按相反顺序从堆中获取元素。

package main

import (
    "container/heap"
    "fmt"
    "math/rand"
)

type IntHeap []int

func (h IntHeap) Len() int           { return len(h) }
func (h IntHeap) Less(i, j int) bool { return h[i] > h[j] }
func (h IntHeap) Swap(i, j int)      { h[i], h[j] = h[j], h[i] }

func (h *IntHeap) Push(x interface{}) {
    *h = append(*h, x.(int))
}

func (h *IntHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[0 : n-1]
    return x
}

const (
    n = 1000000
    m = 5
)

func main() {
    h := &IntHeap{}
    heap.Init(h)

    for i := 0; i < n; i++ {
        x := rand.Intn(n)
        heap.Push(h,x)
        if h.Len() > m {
            heap.Pop(h)
        }
    }

    r := make([]int, h.Len())
    for i := len(r) - 1; i >= 0; i-- {
        r[i] = heap.Pop(h).(int)
    }
    fmt.Printf("%v\n", r)
}

该算法的内存复杂度为M,时间复杂度为N+N*logm+M*logm。