go - Go堆接口(interface)实现的优先级队列大小限制

Question

在 Java 中有一个带有 size 属性的 PriorityQueue。我在这里也有同样的期望(如果我没记错的话)。

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

我正在尝试使用堆接口(interface)来实现这一点。至于我看到的 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。