/prio/prio.go

  1// Copyright 2012 Stefan Nilsson
  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 prio provides a priority queue.
 16// The queue can hold elements that implement the two methods of prio.Interface.
 17package prio
 18
 19/*
 20A type that implements prio.Interface can be inserted into a priority queue.
 21
 22The simplest use case looks like this:
 23
 24	type myInt int
 25
 26	func (x myInt) Less(y prio.Interface) bool { return x < y.(myInt) }
 27	func (x myInt) Index(i int)                {}
 28
 29To use the Remove method you need to keep track of the index of elements
 30in the heap, e.g. like this:
 31
 32	type myType struct {
 33		value int
 34		index int // index in heap
 35	}
 36
 37	func (x *myType) Less(y prio.Interface) bool { return x.value < y.(*myType).value }
 38	func (x *myType) Index(i int)                { x.index = i }
 39*/
 40type Interface interface {
 41	// Less returns whether this element should sort before element x.
 42	Less(x Interface) bool
 43	// Index is called by the priority queue when this element is moved to index i.
 44	Index(i int)
 45}
 46
 47// Queue represents a priority queue.
 48// The zero value for Queue is an empty queue ready to use.
 49type Queue struct {
 50	h []Interface
 51}
 52
 53// New returns an initialized priority queue with the given elements.
 54// A call of the form New(x...) uses the underlying array of x to implement
 55// the queue and hence might change the elements of x.
 56// The complexity is O(n), where n = len(x).
 57func New(x ...Interface) Queue {
 58	q := Queue{x}
 59	heapify(q.h)
 60	return q
 61}
 62
 63// Push pushes the element x onto the queue.
 64// The complexity is O(log(n)), where n = q.Len().
 65func (q *Queue) Push(x Interface) {
 66	n := len(q.h)
 67	q.h = append(q.h, x)
 68	up(q.h, n) // x.Index(n) is done by up.
 69}
 70
 71// Pop removes a minimum element (according to Less) from the queue and returns it.
 72// The complexity is O(log(n)), where n = q.Len().
 73func (q *Queue) Pop() Interface {
 74	h := q.h
 75	n := len(h) - 1
 76	x := h[0]
 77	h[0], h[n] = h[n], nil
 78	h = h[:n]
 79	if n > 0 {
 80		down(h, 0) // h[0].Index(0) is done by down.
 81	}
 82	q.h = h
 83	x.Index(-1) // for safety
 84	return x
 85}
 86
 87// Peek returns, but does not remove, a minimum element (according to Less) of the queue.
 88func (q *Queue) Peek() Interface {
 89	return q.h[0]
 90}
 91
 92// Remove removes the element at index i from the queue and returns it.
 93// The complexity is O(log(n)), where n = q.Len().
 94func (q *Queue) Remove(i int) Interface {
 95	h := q.h
 96	n := len(h) - 1
 97	x := h[i]
 98	h[i], h[n] = h[n], nil
 99	h = h[:n]
100	if i < n {
101		down(h, i) // h[i].Index(i) is done by down.
102		up(h, i)
103	}
104	q.h = h
105	x.Index(-1) // for safety
106	return x
107}
108
109// Len returns the number of elements in the queue.
110func (q *Queue) Len() int {
111	return len(q.h)
112}
113
114// Fix reestablishes the heap ordering after the element at index i has changed its value.
115// Changing the value of the element at index i and then calling Fix is equivalent to,
116// but less expensive than, calling Remove(i) followed by a Push of the new value.
117// The complexity is O(log(n)) where n = q.Len().
118func (q *Queue) Fix(i int) {
119	up(q.h, i)
120	down(q.h, i)
121}
122
123// Establishes the heap invariant in O(n) time.
124func heapify(h []Interface) {
125	n := len(h)
126	for i := n - 1; i >= n/2; i-- {
127		h[i].Index(i)
128	}
129	for i := n/2 - 1; i >= 0; i-- { // h[i].Index(i) is done by down.
130		down(h, i)
131	}
132}
133
134// Moves element at position i towards top of heap to restore invariant.
135func up(h []Interface, i int) {
136	for {
137		parent := (i - 1) / 2
138		if i == 0 || h[parent].Less(h[i]) {
139			h[i].Index(i)
140			break
141		}
142		h[parent], h[i] = h[i], h[parent]
143		h[i].Index(i)
144		i = parent
145	}
146}
147
148// Moves element at position i towards bottom of heap to restore invariant.
149func down(h []Interface, i int) {
150	for {
151		n := len(h)
152		left := 2*i + 1
153		if left >= n || left < 0 { // left < 0 after int overflow
154			h[i].Index(i)
155			break
156		}
157		j := left
158		if right := left + 1; right < n && h[right].Less(h[left]) {
159			j = right
160		}
161		if h[i].Less(h[j]) {
162			h[i].Index(i)
163			break
164		}
165		h[i], h[j] = h[j], h[i]
166		h[i].Index(i)
167		i = j
168	}
169}