golang-heap-prority-queue-max-min-heap

发表于2020-10-14更新于2020-10-14

golang 容器包中 heap 的使用

原文: https://golang.org/pkg/container/heap/

golang 中提供了可以使用 container/heap 操作所有实现了 heap.Interface 类型的 堆数据结构.
操作函数:

1
2
3
4
5
func Fix(h Interface, i int)
func Init(h Interface)
func Pop(h Interface) interface{}
func Push(h Interface, x interface{})
func Remove(h Interface, i int) interface{}

需要实现的接口:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
type Interface interface {
    sort.Interface
    Push(x interface{}) // add x as element Len()
    Pop() interface{}   // remove and return element Len() - 1.
}

type sort.Interface interface {
    // Len is the number of elements in the collection.
    Len() int

    // Less reports whether the element with index i
    // must sort before the element with index j.
    //
    // If both Less(i, j) and Less(j, i) are false,
    // then the elements at index i and j are considered equal.
    // Sort may place equal elements in any order in the final result,
    // while Stable preserves the original input order of equal elements.
    //
    // Less must describe a transitive ordering:
    //  - if both Less(i, j) and Less(j, k) are true, then Less(i, k) must be true as well.
    //  - if both Less(i, j) and Less(j, k) are false, then Less(i, k) must be false as well.
    //
    // Note that floating-point comparison (the < operator on float32 or float64 values)
    // is not a transitive ordering when not-a-number (NaN) values are involved.
    // See Float64Slice.Less for a correct implementation for floating-point values.
    Less(i, j int) bool

    // Swap swaps the elements with indexes i and j.
    Swap(i, j int)
}

堆

  • 最大堆
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
package main

// 堆-完全二叉树

type IntMaxHeap []int

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

func (h *IntMaxHeap) Push(x interface{}) {
    *h = append(*h, x.(int))
}
func (h *IntMaxHeap) Pop() interface{} {
    old := *h
    n := len(old)
    x := old[n-1]
    *h = old[:n-1]
    return x
}

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[:n-1]
    return x
}
  • 最小堆
    修改 Swap(i,j) 函数中的比较逻辑即可

TopK

  • 初始化堆
  • 控制容量
  • 翻转输出
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54

package main

import (
    "container/heap"
)

type TopKHeap []int

type TopK struct {
    TopK    int         // topk
    maxHeap *IntMaxHeap // topk 使用最大堆来实现
    // 先初始化 随机最大堆, 剩下的值 与其中最大的比较,遇到比最大数小 则出队入队操作
}

func NewTopKHeap(nums IntMaxHeap, k int) *TopK {
    t := &TopK{k, &nums}
    heap.Init(t.maxHeap)
    return t
}

func (t *TopK) Push(v interface{}) {
    if v.(int) < (*t.maxHeap)[0] {
        heap.Pop(t.maxHeap)
        heap.Push(t.maxHeap, v)
    }
}

func (t *TopK) Pop() interface{} {
    return heap.Pop(t.maxHeap)
}

func (t *TopK) Len() int {
    return t.maxHeap.Len()
}

func (t *TopK) Max() interface{} {
    return (*t.maxHeap)[0]
}

// 翻转
func (t *TopK) Reversal() []int {
    l := t.maxHeap.Len()
    r := make([]int, l)
    var n IntMaxHeap = *t.maxHeap
    for l > 0 {
        l--
        r[l] = t.Pop().(int)
    }
    t.maxHeap = &n
    heap.Init(t.maxHeap)
    return r
}

优先级队列

  • 消息结构化
  • 优先级字段
  • 更新优先级
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
package main

import (
    "container/heap"
)

// 堆结构的优先级队列

// 消息体
type Item struct {
    value    string
    priority int
    index    int
}

// 堆
type PriorityQueue []*Item

func (p PriorityQueue) Len() int           { return len(p) }
func (p PriorityQueue) Less(i, j int) bool { return p[i].priority > p[j].priority }
func (p PriorityQueue) Swap(i, j int) {
    p[i], p[j] = p[j], p[i]
    p[i].index = i
    p[j].index = j
}

func (p *PriorityQueue) Push(x interface{}) {
    n := len(*p)
    item := x.(*Item)
    item.index = n
    *p = append(*p, item)
}

func (p *PriorityQueue) Pop() interface{} {
    old := *p
    n := len(old)
    item := old[n-1]
    item.index = -1
    *p = old[:n-1]
    return item
}

func (p *PriorityQueue) Update(item *Item, value string, priority int) {
    item.value = value
    item.priority = priority
    heap.Fix(p, item.index)
}

测试

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
package main

import (
    "container/heap"
    "fmt"
    "testing"
)

func TestIntHeap(t *testing.T) {
    h := &IntHeap{3, 4, 6, 7, 2, 1}
    heap.Init(h)
    heap.Push(h, 9)
    fmt.Printf("mininum: %d\n", heap.Pop(h))
    heap.Push(h, 1)
    for h.Len() > 0 {
        fmt.Printf("%d ", heap.Pop(h))
    }
}

func TestTopKHeap(t *testing.T) {
    i, k := 4, 4
    h := &IntMaxHeap{9, 4, 6, 7, 2, 1}
    l := len(*h)
    initNum := (*h)[:k]
    topk := NewTopKHeap(initNum, k)

    for i < l {
        topk.Push((*h)[i])
        i++
    }
    topk.Push(8)
    topk.Push(8)
    fmt.Printf("topk's maxnum: %d\n", topk.Max())
    fmt.Println(topk.Reversal())
    for topk.Len() > 0 {
        fmt.Printf("%d ", topk.Pop())
    }
}

func TestHepPQ(t *testing.T) {
    items := map[string]int{"taskA": 1, "taskB": 2, "taskC": 3}
    p := make(PriorityQueue, len(items))
    i := 0
    for value, priority := range items {
        p[i] = &Item{value: value, priority: priority, index: i}
        i++
    }
    heap.Init(&p)
    item := &Item{
        value: "ooo", priority: 4,
    }
    heap.Push(&p, item)
    p.Update(item, item.value, 5)
    for p.Len() > 0 {
        item := heap.Pop(&p).(*Item)
        fmt.Printf("%.2d:%s\n", item.priority, item.value)
    }
}

输出结果

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
oomcc:gogin oom$ go test -timeout 30s -run ^TestIntHeap$ ./src/learn/heap*.go -v
=== RUN   TestIntHeap
mininum: 1
1 2 3 4 6 7 9 --- PASS: TestIntHeap (0.00s)
PASS
ok      command-line-arguments  0.007s
oomcc:gogin oom$ go test -timeout 30s -run ^TestTopKHeap$ ./src/learn/heap*.go -v
=== RUN   TestTopKHeap
topk s maxnum: 6
[1 2 4 6]
6 4 2 1 --- PASS: TestTopKHeap (0.00s)
PASS
ok      command-line-arguments  0.007s
oomcc:gogin oom$ go test -timeout 30s -run ^TestHepPQ$ ./src/learn/heap*.go -v 
=== RUN   TestHepPQ
05:ooo
03:taskC
02:taskB
01:taskA
--- PASS: TestHepPQ (0.00s)
PASS
ok      command-line-arguments  0.006s