1// Copyright 2012 The Go Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5// This example demonstrates a priority queue built using the heap interface.
6package heap_test
7
8import (
9	"container/heap"
10	"fmt"
11)
12
13// An Item is something we manage in a priority queue.
14type Item struct {
15	value    string // The value of the item; arbitrary.
16	priority int    // The priority of the item in the queue.
17	// The index is needed by update and is maintained by the heap.Interface methods.
18	index int // The index of the item in the heap.
19}
20
21// A PriorityQueue implements heap.Interface and holds Items.
22type PriorityQueue []*Item
23
24func (pq PriorityQueue) Len() int { return len(pq) }
25
26func (pq PriorityQueue) Less(i, j int) bool {
27	// We want Pop to give us the highest, not lowest, priority so we use greater than here.
28	return pq[i].priority > pq[j].priority
29}
30
31func (pq PriorityQueue) Swap(i, j int) {
32	pq[i], pq[j] = pq[j], pq[i]
33	pq[i].index = i
34	pq[j].index = j
35}
36
37func (pq *PriorityQueue) Push(x any) {
38	n := len(*pq)
39	item := x.(*Item)
40	item.index = n
41	*pq = append(*pq, item)
42}
43
44func (pq *PriorityQueue) Pop() any {
45	old := *pq
46	n := len(old)
47	item := old[n-1]
48	old[n-1] = nil  // don't stop the GC from reclaiming the item eventually
49	item.index = -1 // for safety
50	*pq = old[0 : n-1]
51	return item
52}
53
54// update modifies the priority and value of an Item in the queue.
55func (pq *PriorityQueue) update(item *Item, value string, priority int) {
56	item.value = value
57	item.priority = priority
58	heap.Fix(pq, item.index)
59}
60
61// This example creates a PriorityQueue with some items, adds and manipulates an item,
62// and then removes the items in priority order.
63func Example_priorityQueue() {
64	// Some items and their priorities.
65	items := map[string]int{
66		"banana": 3, "apple": 2, "pear": 4,
67	}
68
69	// Create a priority queue, put the items in it, and
70	// establish the priority queue (heap) invariants.
71	pq := make(PriorityQueue, len(items))
72	i := 0
73	for value, priority := range items {
74		pq[i] = &Item{
75			value:    value,
76			priority: priority,
77			index:    i,
78		}
79		i++
80	}
81	heap.Init(&pq)
82
83	// Insert a new item and then modify its priority.
84	item := &Item{
85		value:    "orange",
86		priority: 1,
87	}
88	heap.Push(&pq, item)
89	pq.update(item, item.value, 5)
90
91	// Take the items out; they arrive in decreasing priority order.
92	for pq.Len() > 0 {
93		item := heap.Pop(&pq).(*Item)
94		fmt.Printf("%.2d:%s ", item.priority, item.value)
95	}
96	// Output:
97	// 05:orange 04:pear 03:banana 02:apple
98}
99