Implementing a Priority Queue Using Heap in Go

In computer science, a priority queue is an abstract data type that stores elements with priorities. It allows elements with higher priorities to be dequeued first. Priority queues are widely used in many algorithms such as Dijkstra’s shortest path algorithm, Prim’s minimum spanning tree algorithm, Huffman coding, and many more. In this blog post, we will explore how to implement a priority queue using a heap in the Go programming language.

A heap is a tree-based data structure that satisfies the heap property. In a heap, every parent node has a priority greater than or equal to its children. There are two types of heaps: max heaps and min heaps. In a max heap, the root node has the highest priority, and in a min heap, the root node has the lowest priority.

A heap can be represented using an array. The root node is stored in the first index of the array, and its left and right children are stored in the second and third indices, respectively. The left child of a node at index i is located at index 2i, and the right child is located at index 2i+1. The parent of a node at index i is located at index i/2.

In Go, we can implement a priority queue using the container/heap package. The package provides an interface for heap operations such as pushing, popping, and updating elements. To use the container/heap package, we need to implement the heap.Interface interface.

The heap.Interface interface requires three methods to be implemented: Len(), Less(), and Swap(). Len() returns the length of the heap, Less() returns true if the element at index i has a higher priority than the element at index j, and Swap() swaps the elements at indices i and j.

Here is an implementation of a priority queue using a max heap in Go:

package main

import (
	"container/heap"
	"fmt"
)

// An Item is something we manage in a priority queue.
type Item struct {
	value    string // The value of the item; arbitrary.
	priority int    // The priority of the item in the queue.
	// The index is needed by update and is maintained by the heap.Interface methods.
	index int // The index of the item in the heap.
}

// A PriorityQueue implements heap.Interface and holds Items.
type PriorityQueue []*Item

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

func (pq PriorityQueue) Less(i, j int) bool {
	// We want Pop to give us the highest, not lowest, priority so we use greater than here.
	return pq[i].priority > pq[j].priority
}

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

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

func (pq *PriorityQueue) Pop() any {
	old := *pq
	n := len(old)
	item := old[n-1]
	old[n-1] = nil  // avoid memory leak
	item.index = -1 // for safety
	*pq = old[0 : n-1]
	return item
}

// 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)
}

// This example creates a PriorityQueue with some items, adds and manipulates an item,
// and then removes the items in priority order.
func main() {
	// Some items and their priorities.
	items := map[string]int{
		"banana": 3, "apple": 2, "pear": 4,
	}

	// Create a priority queue, put the items in it, and
	// establish the priority queue (heap) invariants.
	pq := make(PriorityQueue, len(items))
	i := 0
	for value, priority := range items {
		pq[i] = &Item{
			value:    value,
			priority: priority,
			index:    i,
		}
		i++
	}
	heap.Init(&pq)

	// Insert a new item and then modify its priority.
	item := &Item{
		value:    "orange",
		priority: 1,
	}
	heap.Push(&pq, item)
	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 \n", item.priority, item.value)
	}
}

/*
Output:

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

*/

Implementing a priority queue using a heap in Go is straightforward using the container/heap package. Priority queues are useful in many algorithms, and understanding how to implement them using a heap can improve the efficiency of your programs. With the example above, you should now have a good understanding of how to use a priority queue in your Go programs.