Table of Contents

• Implement PriorityQueue

Implement PriorityQueue

Note how PriorityQueue manage indexes of elements internally, to utilize APIs like heap.Fix(), and heap.Remove(), etc.

package main

import (
    "container/heap"
    "fmt"
)

type Item struct {
    value    string
    priority int
    index    int
}

type PriorityQueue []*Item

// Methods for sort.Interface
func (pq PriorityQueue) Len() int {
    return len(pq)
}

func (pq PriorityQueue) Less(i, j int) bool {
    // Max Heap
    return pq[i].priority > pq[j].priority
}

func (pq PriorityQueue) Swap(i, j int) {
    pq[i], pq[j] = pq[j], pq[i]
    // Note that the items keep the indexes of themselves.
    pq[i].index = i
    pq[j].index = j
}

// Interface.Push(x interface{}) should add x as element Len()
func (pq *PriorityQueue) Push(x interface{}) {
    n := len(*pq)
    item := x.(*Item)
    item.index = n
    *pq = append(*pq, item)
}

// Interface.Pop() interface{} should remove and return element Len() - 1
func (pq *PriorityQueue) Pop() interface{} {
    old := *pq
    n := len(old)
    item := old[n-1]
    item.index = -1 // for safety
    *pq = old[0 : n-1]
    return item
}

// Custom operations:

// update modifies the priority and value of an Item in the queue.
func (pq *PriorityQueue) update(item *Item, value string, priority int) {
    item.value = value
    item.priority = priority
    heap.Fix(pq, item.index) // Re-heap the item.
}

func main() {
    items := map[string]int{"banana": 3, "apple": 2, "pear": 4}

    // Create a priority queue, put the items in it.
    pq := make(PriorityQueue, len(items))

    i := 0
    for value, priority := range items {
        pq[i] = &Item{
            value:    value,
            priority: priority,
            index:    i,
        }
        i++
    }

    // Establish the priority queue (heap) invariants.
    heap.Init(&pq)

    // Insert a new item and then modify its priority.
    item := &Item{
        value:    "orange",
        priority: 1,
    }
    heap.Push(&pq, item)

    // Update the just pushed item's priority.
    pq.update(item, item.value, 5)

    // Take the items out; they arrive in decreasing priority order.
    for pq.Len() > 0 {
        item := heap.Pop(&pq).(*Item)
        fmt.Printf("%.2d:%s ", item.priority, item.value)
    }
}

05:orange 04:pear 03:banana 02:apple