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

Question

In Java there is a PriorityQueue with size attribute. I'm expecting the same here (if I'm not wrong).

Use case: Read millions of data one by one and send it to a priority queue. I want only the top 5 calculated elements so I want a heap/priority queue of size 5 only.

I'm trying to use heap interface to achieve this. As to what I see golang does dynamic array increase but this is not feasible in my case.

I am referring to https://play.golang.org/p/wE413xwmxE

How can I achieve this?

Answer 1

If you just want top M smallest elements from N elements then use the heap that would give you the biggest element and prune the heap whenever its size goes over value M. Then take elements from the heap in reverse order.

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
", r)
}

This algorithm has memory complexity of M and time complexity of N + N * log M + M * log M.