github.com/code-reading/golang@v0.0.0-20220303082512-ba5bc0e589a3/coding/container/00-heap-pq.go (about)

     1  package main
     2  
     3  import (
     4  	"container/heap"
     5  	"fmt"
     6  )
     7  
     8  // Item 优先级队列项Item
     9  type Item struct {
    10  	value    string // 数据
    11  	priority int    // 优先级
    12  	index    int    // 堆数组中的索引
    13  }
    14  
    15  // PriorityQueue  优先级队列
    16  type PriorityQueue []*Item
    17  
    18  // 接着heap 接口实现优先级队列， 需要实现 heap 的Push 和Pop
    19  // 及其继承sort.Interface 的 Len, Less, Swap 这五个接口
    20  
    21  // Len 不涉及修改pq 操作的 可以不需要使用指针
    22  func (p PriorityQueue) Len() int {
    23  	return len(p)
    24  }
    25  
    26  // Less 业务定义: 优先级大的 表示优先级高 这里用 > 大于号
    27  func (p PriorityQueue) Less(i, j int) bool {
    28  	return p[i].priority > p[j].priority
    29  }
    30  
    31  // Swap 交换i和j的位置
    32  func (p PriorityQueue) Swap(i, j int) {
    33  	p[i], p[j] = p[j], p[i]
    34  	p[i].index = i
    35  	p[j].index = j
    36  }
    37  
    38  // 涉及到修改堆数组的时候需要用指针
    39  // Push 加入新节点, 注意 堆接口Push 的入口参数是interface{}
    40  func (p *PriorityQueue) Push(x interface{}) {
    41  	// 获得加入新节点的下标位置
    42  	n := len(*p)
    43  	item := x.(*Item)
    44  	item.index = n          // 保存数组索引
    45  	(*p) = append(*p, item) // 添加item
    46  }
    47  
    48  // Pop 出口参数也是interface
    49  func (p *PriorityQueue) Pop() interface{} {
    50  	old := *p        // 用临时变量old 指向优先级队列p
    51  	n := len(old)    // 获取队列长度
    52  	item := old[n-1] // 获取最后一个参数
    53  	// 注意事项
    54  	// 防止内存泄露
    55  	old[n-1] = nil
    56  	// 防止索引被引用
    57  	item.index = -1
    58  	*p = old[:n-1] // 修改队列长度
    59  	return item    // 返回堆顶元素
    60  }
    61  
    62  // update  更新索引和值
    63  func (p *PriorityQueue) update(item *Item, value string, priority int) {
    64  	item.value = value
    65  	item.priority = priority
    66  	heap.Fix(p, item.index) // 索引在更新堆化这里用到, 因为重新调整堆时 需要知道是哪个节点要调整
    67  }
    68  
    69  func main() {
    70  	items := map[string]int{
    71  		"banana": 3, "apple": 2, "pear": 4,
    72  	}
    73  	pq := make(PriorityQueue, len(items)) // 初始化优先级队列, 注意这里填写了len 后面就不要用append 了否则前面len个数据是Nil
    74  	i := 0
    75  	for item, priority := range items {
    76  		pq[i] = &Item{
    77  			value:    item,
    78  			priority: priority,
    79  			index:    i,
    80  		}
    81  		i++
    82  	}
    83  	heap.Init(&pq) // 构建最小堆
    84  	// 插入一个新项，并修改其优先级
    85  	item := &Item{
    86  		value:    "orange",
    87  		priority: 1,
    88  	}
    89  	heap.Push(&pq, item)
    90  	pq.update(item, item.value, 5) // 修改优先级
    91  
    92  	// 读取优先级队列
    93  	for pq.Len() > 0 {
    94  		item := heap.Pop(&pq).(*Item)
    95  		fmt.Printf("%.2d:%s ", item.priority, item.value)
    96  	}
    97  	// 05:orange 04:pear 03:banana 02:apple
    98  }