Leaky Cauldron
0%

Golang堆的实现

Posted on Edited on

堆的使用

Golang 中没有给出堆的具体实现，仅提供了标准容器库 container/heap ，需手动实现。

heap 接口定义如下

1
2
3
4
5
type Interface interface {
	sort.Interface
	Push(x interface{}) // add x as element Len()
	Pop() interface{}   // remove and return element Len() - 1.
}

该接口同时继承了 sort.Interface ，定义如下

1
2
3
4
5
6
7
8
9
type Interface interface {
	// Len is the number of elements in the collection.
	Len() int
	// Less reports whether the element with
	// index i should sort before the element with index j.
	Less(i, j int) bool
	// Swap swaps the elements with indexes i and j.
	Swap(i, j int)
}

故创建堆时需实现 Len() 、Less(i, j int) 、Swap(i, j int) 、Push(x interface{}) 、Pop() 五个方法。

向堆中插入节点时，堆会自动调整结构，Golang 中始终保持最后一个元素为堆顶，故 Pop 时返回最后一个节点并删除即可。

实例

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
package main

import (
	"container/heap"
	"fmt"
	"math/rand"
	"strconv"
	"time"
)

type Person struct {
	Name string
	Age  int
}

type PersonHeap []Person

func (ph PersonHeap) Len() int {
	return len(ph)
}

func (ph PersonHeap) Swap(i, j int) {
	ph[i], ph[j] = ph[j], ph[i]
}

func (ph PersonHeap) Less(i, j int) bool {
	return ph[i].Age < ph[j].Age // 大根堆
}

func (ph *PersonHeap) Push(p interface{}) {
	*ph = append(*ph, p.(Person))
}

func (ph *PersonHeap) Pop() interface{} {
	n := len(*ph)
	result := (*ph)[n-1]
	*ph = (*ph)[:n-1]
	return result
}

func main() {
	rand.Seed(time.Now().Unix())
	hp := &PersonHeap{}
	for i := 0; i < 6; i++ {
		*hp = append(*hp, Person{"p" + strconv.Itoa(i), rand.Intn(50)})
	}
	fmt.Printf("原始切片: %#v\n", hp)

	heap.Init(hp)

	fmt.Printf("堆: %#v\n", hp)

	fmt.Println(hp.Pop())
	heap.Push(hp, Person{"newP", 30})

	fmt.Printf("%#v\n", hp)
}

例题(LeetCode #23)

合并 k 个排序链表，返回合并后的排序链表。请分析和描述算法的复杂度。

示例:

输入:
[
  1->4->5,
  1->3->4,
  2->6
]
输出: 1->1->2->3->4->4->5->6

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
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
package main

import (
	"container/heap"
	"fmt"
)

type ListNode struct {
	Val  int
	Next *ListNode
}

type PriorityQueue []*ListNode

func (pq PriorityQueue) Len() int {
	return len(pq)
}

func (pq PriorityQueue) Less(i, j int) bool {
	return pq[i].Val < pq[j].Val
}

func (pq *PriorityQueue) Swap(i, j int) {
	(*pq)[i].Val, (*pq)[j].Val = (*pq)[j].Val, (*pq)[i].Val
}

func (pq *PriorityQueue) Push(x interface{}) {
	*pq = append(*pq, x.(*ListNode))
}

func (pq *PriorityQueue) Pop() interface{} {
	result := (*pq)[len(*pq)-1]
	*pq = (*pq)[:len(*pq)-1]
	return result
}

func mergeKLists(lists []*ListNode) *ListNode {
	if lists == nil || len(lists) == 0 {
		return nil
	}

	dummyHead := &ListNode{0, nil}
	curNode := dummyHead
	var pq PriorityQueue

	for _, list := range lists {
		tmp := list
		for tmp != nil {
			heap.Push(&pq, tmp)
			tmp = tmp.Next
		}
	}

	for len(pq) != 0 {
		curNode.Next = heap.Pop(&pq).(*ListNode)
		curNode = curNode.Next
	}
	curNode.Next = nil

	return dummyHead.Next
}

func main() {
	l1 := &ListNode{1, nil}
	l1.Next = &ListNode{4, nil}
	l1.Next.Next = &ListNode{5, nil}

	l2 := &ListNode{1, nil}
	l2.Next = &ListNode{3, nil}
	l2.Next.Next = &ListNode{4, nil}

	l3 := &ListNode{2, nil}
	l3.Next = &ListNode{6, nil}

	lists := []*ListNode{l1, l2, l3}

	result := mergeKLists(lists)
	for result != nil {
		fmt.Printf("%d ", result.Val)
		result = result.Next
	}
}