priority_queue.go

package main

import "container/heap"

type Item struct {
	value    interface{}
	priority int
	index    int
}

type PriorityQueue []*Item

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

func (pq PriorityQueue) Less(i, j int) bool {
	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 interface{}) {
	n := len(*pq)
	item := x.(*Item)
	item.index = n
	*pq = append(*pq, item)
}

func (pq *PriorityQueue) Pop() interface{} {
	old := *pq
	n := len(old)
	item := old[n-1]
	item.index = -1
	*pq = old[:n-1]
	return item
}

type PriorityQueueWrapper struct {
	pq    PriorityQueue
	items map[interface{}]*Item
}

func NewPriorityQueueWrapper() *PriorityQueueWrapper {
	return &PriorityQueueWrapper{
		pq:    make(PriorityQueue, 0),
		items: make(map[interface{}]*Item),
	}
}

func (pw *PriorityQueueWrapper) Enqueue(value interface{}, priority int) {
	item := &Item{
		value:    value,
		priority: priority,
	}
	heap.Push(&(pw.pq), item)
	pw.items[value] = item
}

func (pw *PriorityQueueWrapper) Dequeue() interface{} {
	item := heap.Pop(&(pw.pq)).(*Item)
	delete(pw.items, item.value)
	return item.value
}

func (pw *PriorityQueueWrapper) IsEmpty() bool {
	return pw.pq.Len() == 0
}

func (pw *PriorityQueueWrapper) Size() int {
	return pw.pq.Len()
}