Go 堆（优先队列）的两种实现方式

堆是一棵完全二叉树（一种平衡二叉树）

最小堆：父节点的值小于左右儿子节点
最大堆：父节点的值大于左右儿子节点

手搓

初始化一个空数组 arr
向堆中添加一个元素 item 时：
1. 在堆底添加 append(arr, item)
2. 如果该元素 i = len(arr) - 1 比其父节点 (i - 1) / 2 大，则将二者交换
3. 不断向上交换，直至根节点（堆顶）或者其父节点比该元素大
删除堆顶（最大）元素时：
1. 把堆底元素放到根节点（堆顶）上 arr[0] = arr[len(arr) - 1]，并删除该元素 arr = arr[:len(arr) - 1]
2. 如果该元素 i = 0 比其较大的儿子节点 2*i + 1 / 2*i + 2 小，则将该元素和较大的儿子节点交换
3. 不断向下交换，直至堆底或者其较大的儿子节点比该元素小

type Item struct {
    Value    int
    Priority int
}

type Heap struct {
    arr []*Item
}

func (h *Heap) Push(item *Item) {
    h.arr = append(h.arr, item)

    for i := len(h.arr) - 1; i > 0; {
       parent := (i - 1) / 2

       if item.Priority >= h.arr[parent].Priority {
          break
       }

       h.arr[i], h.arr[parent] = h.arr[parent], h.arr[i]
       i = parent
    }
}

func (h *Heap) Pop() *Item {
    if len(h.arr) == 0 {
       panic("empty")
    }

    ret := h.arr[0]

    h.arr[0] = h.arr[len(h.arr)-1]
    h.arr = h.arr[:len(h.arr)-1]

    for i := 0; len(h.arr) > 0; {
       a := 2*i + 1    // 左儿子节点
       b := 2*i + 2    // 右儿子节点

       if a >= len(h.arr) {
          break
       }

       if b < len(h.arr) && h.arr[a].Priority > h.arr[b].Priority {
          a = b
       }

       if h.arr[i].Priority <= h.arr[a].Priority {
          break
       }

       h.arr[i], h.arr[a] = h.arr[a], h.arr[i]
       i = a
    }

    return ret
}

Go SDK

Go SDK 提供了 container/heap 包来实现堆，container/heap 包定义了一个接口 heap.Interface，该接口包含以下方法：

Len() int：返回堆的元素数量
Less(i, j int) bool：判断索引为i的元素是否小于索引为j的元素
Swap(i, j int)：交换索引为i和j的元素
Push(x interface{})：向堆中添加一个元素
Pop() interface{}：从堆中删除并返回堆顶元素

type Item struct {
    Value    int
    Priority int
}

type PriorityQueue []*Item

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

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

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

func (pq *PriorityQueue) Push(item interface{}) {
    *pq = append(*pq, item.(*Item))
}

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

    return item
}

func (pq *PriorityQueue) Enqueue(value int, priority int) {
    heap.Push(pq, &Item{
       Value:    value,
       Priority: priority,
    })
}

func (pq *PriorityQueue) Dequeue() int {
    item := heap.Pop(pq).(*Item)
    return item.Value
}

手搓

Go SDK

发送评论 编辑评论

推荐文章

发送评论编辑评论