1// Copyright 2017 Google Inc. 2// 3// Licensed under the Apache License, Version 2.0 (the "License"); 4// you may not use this file except in compliance with the License. 5// You may obtain a copy of the License at 6// 7// http://www.apache.org/licenses/LICENSE-2.0 8// 9// Unless required by applicable law or agreed to in writing, software 10// distributed under the License is distributed on an "AS IS" BASIS, 11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12// See the License for the specific language governing permissions and 13// limitations under the License. 14 15// Package pq provides a priority queue. 16package pq 17 18import "container/heap" 19 20// NewQueue returns an unbounded priority queue that compares elements using 21// less; the minimal element is at the top of the queue. 22// 23// If setIndex is not nil, the queue calls setIndex to inform each element of 24// its position in the queue. If an element's priority changes, its position in 25// the queue may be incorrect. Call Fix on the element's index to update the 26// queue. Call Remove on the element's index to remove it from the queue. 27func NewQueue(less func(x, y interface{}) bool, setIndex func(x interface{}, idx int)) *Queue { 28 return &Queue{ 29 heap: pqHeap{ 30 less: less, 31 setIndex: setIndex, 32 }, 33 } 34} 35 36// Queue is a priority queue that supports updating the priority of an element. 37// A Queue must be created with NewQueue. 38type Queue struct { 39 heap pqHeap 40} 41 42// Len returns the number of elements in the queue. 43func (pq *Queue) Len() int { 44 return pq.heap.Len() 45} 46 47// Push adds x to the queue. 48func (pq *Queue) Push(x interface{}) { 49 heap.Push(&pq.heap, x) 50} 51 52// Min returns the minimal element. 53// Min panics if the queue is empty. 54func (pq *Queue) Min() interface{} { 55 return pq.heap.a[0] 56} 57 58// Pop removes and returns the minimal element. 59// Pop panics if the queue is empty. 60func (pq *Queue) Pop() interface{} { 61 return heap.Pop(&pq.heap) 62} 63 64// Fix adjusts the heap to reflect that the element at index has changed priority. 65func (pq *Queue) Fix(index int) { 66 heap.Fix(&pq.heap, index) 67} 68 69// Remove removes the element at index i from the heap. 70func (pq *Queue) Remove(index int) { 71 heap.Remove(&pq.heap, index) 72} 73 74// pqHeap implements heap.Interface. 75type pqHeap struct { 76 a []interface{} 77 less func(x, y interface{}) bool 78 setIndex func(x interface{}, idx int) 79} 80 81func (h pqHeap) Len() int { 82 return len(h.a) 83} 84 85func (h pqHeap) Less(i, j int) bool { 86 return h.less(h.a[i], h.a[j]) 87} 88 89func (h pqHeap) Swap(i, j int) { 90 h.a[i], h.a[j] = h.a[j], h.a[i] 91 if h.setIndex != nil { 92 h.setIndex(h.a[i], i) 93 h.setIndex(h.a[j], j) 94 } 95} 96 97func (h *pqHeap) Push(x interface{}) { 98 n := len(h.a) 99 if h.setIndex != nil { 100 h.setIndex(x, n) 101 } 102 h.a = append(h.a, x) 103} 104 105func (h *pqHeap) Pop() interface{} { 106 old := h.a 107 n := len(old) 108 x := old[n-1] 109 h.a = old[:n-1] 110 return x 111} 112