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

Question

在Java中，有一个带有size属性的PriorityQueue。我在这里期待同样的事情（如果我没错的话）。

使用案例：逐个读取数百万个数据并将其发送到优先级队列。我只想要前5个计算元素，所以我只想要一个大小为5的堆/优先级队列。

我正在尝试使用堆接口来实现这一目标。至于我看到golang的动态数组增加，但在我的情况下这是不可行的。

我指的是https://play.golang.org/p/wE413xwmxE

我怎样才能做到这一点？

Answer 1

如果您只想要来自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 * log M + M * log M的时间复杂度。