Priority Queue Implementation in Go

Question

I have just seen an implementation of a priority queue in a generic kind of way in which any type satisfying an interface can be put into the queue. Is this the way to go with go or does this introduces any issues?

// Copyright 2012 Stefan Nilsson
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package prio provides a priority queue.
// The queue can hold elements that implement the two methods of prio.Interface.
package prio

/*
A type that implements prio.Interface can be inserted into a priority queue.

The simplest use case looks like this:

        type myInt int

        func (x myInt) Less(y prio.Interface) bool { return x < y.(myInt) }
        func (x myInt) Index(i int)                {}

To use the Remove method you need to keep track of the index of elements
in the heap, e.g. like this:

        type myType struct {
                value int
                index int // index in heap
        }

        func (x *myType) Less(y prio.Interface) bool { return x.value < y.(*myType).value }
        func (x *myType) Index(i int)                { x.index = i }
*/
type Interface interface {
        // Less returns whether this element should sort before element x.
        Less(x Interface) bool
        // Index is called by the priority queue when this element is moved to index i.
        Index(i int)
}

// Queue represents a priority queue.
// The zero value for Queue is an empty queue ready to use.
type Queue struct {
        h []Interface
}

// New returns an initialized priority queue with the given elements.
// A call of the form New(x...) uses the underlying array of x to implement
// the queue and hence might change the elements of x.
// The complexity is O(n), where n = len(x).
func New(x ...Interface) Queue {
        q := Queue{x}
        heapify(q.h)
        return q
}

// Push pushes the element x onto the queue.
// The complexity is O(log(n)) where n = q.Len().
func (q *Queue) Push(x Interface) {
        n := len(q.h)
        q.h = append(q.h, x)
        up(q.h, n) // x.Index(n) is done by up.
}

// Pop removes a minimum element (according to Less) from the queue and returns it.
// The complexity is O(log(n)), where n = q.Len().
func (q *Queue) Pop() Interface {
        h := q.h
        n := len(h) - 1
        x := h[0]
        h[0], h[n] = h[n], nil
        h = h[:n]
        if n > 0 {
                down(h, 0) // h[0].Index(0) is done by down.
        }
        q.h = h
        x.Index(-1) // for safety
        return x
}

// Peek returns, but does not remove, a minimum element (according to Less) of the queue.
func (q *Queue) Peek() Interface {
        return q.h[0]
}

// Remove removes the element at index i from the queue and returns it.
// The complexity is O(log(n)), where n = q.Len().
func (q *Queue) Remove(i int) Interface {
        h := q.h
        n := len(h) - 1
        x := h[i]
        h[i], h[n] = h[n], nil
        h = h[:n]
        if i < n {
                down(h, i) // h[i].Index(i) is done by down.
                up(h, i)
        }
        q.h = h
        x.Index(-1) // for safety
        return x
}

// Len returns the number of elements in the queue.
func (q *Queue) Len() int {
        return len(q.h)
}

// Establishes the heap invariant in O(n) time.
func heapify(h []Interface) {
        n := len(h)
        for i := n - 1; i >= n/2; i-- {
                h[i].Index(i)
        }
        for i := n/2 - 1; i >= 0; i-- { // h[i].Index(i) is done by down.
                down(h, i)
        }
}

// Moves element at position i towards top of heap to restore invariant.
func up(h []Interface, i int) {
        for {
                parent := (i - 1) / 2
                if i == 0 || h[parent].Less(h[i]) {
                        h[i].Index(i)
                        break
                }
                h[parent], h[i] = h[i], h[parent]
                h[i].Index(i)
                i = parent
        }
}

// Moves element at position i towards bottom of heap to restore invariant.
func down(h []Interface, i int) {
        for {
                n := len(h)
                left := 2*i + 1
                if left >= n {
                        h[i].Index(i)
                        break
                }
                j := left
                if right := left + 1; right < n && h[right].Less(h[left]) {
                        j = right
                }
                if h[i].Less(h[j]) {
                        h[i].Index(i)
                        break
                }
                h[i], h[j] = h[j], h[i]
                h[i].Index(i)
                i = j
        }
}

Priority Queue Implementation in Go

Answers (1)

Related Questions