使用 go 实现优先队列

问题

输入整数数组 arr ,找出其中最小的 k 个数。例如,输入4、5、1、6、2、7、3、8这8个数字,则最小的4个数字是1、2、3、4。

代码

注意看看,用 go 实现堆是如何实现的?

package main

import (
	"container/heap"
)

type IntHeap []int

func (h IntHeap) Len() int { return len(h) }

// 为了实现大根堆,Less在大于时返回小于
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] }

// Push and Pop use pointer receivers because they modify the slice's length,
// not just its contents.
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
}

// 大根堆
func getLeastNumbers(arr []int, k int) []int {
	if k == 0 {
		return []int{}
	}
	h := make(IntHeap, k)
	hp := &h
	copy(h, IntHeap(arr[:k+1]))
	heap.Init(hp)
	for i := k; i < len(arr); i++ {
		if arr[i] < h[0] {
			heap.Pop(hp)
			heap.Push(hp, arr[i])
		}
	}
	return h
}

优先队列2

package main

import (
    "container/heap"
    "fmt"
)

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
}

func main() {
    h := &IntHeap{2, 1, 5, 100, 3, 6, 4, 5}
    heap.Init(h)
    heap.Push(h, 3)
    fmt.Printf("minimum: %d\n", (*h)[0])
    for h.Len() > 0 {
        fmt.Printf("%d ", heap.Pop(h))
    }
}
posted @ 2021-08-27 21:45  沧海一声笑rush  阅读(275)  评论(0编辑  收藏  举报