Indexed priority queue

heap.go

package heap

import (
	"math"
)

type Interface interface {
	Len() int
	Less(i, j int) bool
	Swap(i, j int)
	Push(v interface{})
	Pop() interface{}
}

func parent(i int) int {
	return int(math.Floor(float64(i-1) / 2))
}

func down(h Interface, i int) {
	for i < h.Len() {
		l := 2*i + 1
		r := 2*i + 2

		if l >= h.Len() {
			break
		}

		if r >= h.Len() {
			r = l
		}

		min, max := r, l

		if h.Less(l, r) {
			min, max = l, r
		}

		if h.Less(min, i) {
			h.Swap(min, i)
			i = min
			continue
		}

		if h.Less(max, i) {
			h.Swap(max, i)
			i = max
			continue
		}

		break
	}
}

func up(h Interface, i int) {
	for i > 0 {
		p := parent(i)

		if h.Less(p, i) {
			break
		}

		h.Swap(p, i)
		i = p
	}
}

func Push(h Interface, v interface{}) {
	h.Push(v)
	up(h, h.Len()-1)
}

func Pop(h Interface) interface{} {
	h.Swap(0, h.Len()-1)
	v := h.Pop()
	down(h, 0)
	return v
}

func Fix(h Interface, i int) {
	p := parent(i)
	if h.Less(i, p) {
		up(h, i)
	} else {
		down(h, i)
	}
}

func Del(h Interface, i int) {
	h.Swap(i, h.Len() - 1)
	h.Pop()
	down(h, i)
}

heap_test.go

package heap_test

import (
	myheap "heap"
	"container/heap"
	"testing"
)

type IntHeap []int

func (h IntHeap) Len() int {
	return len(h)
}

func (h IntHeap) Less(i, j int) bool {
	return h[i] < h[j]
}

func (h IntHeap) Swap(i, j int) {
	h[i], h[j] = h[j], h[i]
}

func (h *IntHeap) Push(x interface{}) {
	*h = append(*h, x.(int))
}

func (h *IntHeap) Pop() interface{} {
	old := *h
	n := len(old)
	x := old[n-1]
	*h = old[0 : n-1]
	return x
}

func TestHeap(t *testing.T) {
	myh := &IntHeap{}
	myheap.Push(myh, 10)
	myheap.Push(myh, 8)
	myheap.Push(myh, 11)
	myheap.Push(myh, 5)
	myheap.Push(myh, 4)
	myheap.Push(myh, 3)
	myheap.Push(myh, 20)
	t.Log(myh)

	/* t.Log(myheap.Pop(myh))
	t.Log(myheap.Pop(myh))
	t.Log(myheap.Pop(myh))
	t.Log(myh) */
	[]int(*myh)[3] = 1
	myheap.Fix(myh, 3)
	[]int(*myh)[2] = 12
	myheap.Fix(myh, 2)
	t.Log(myh)

	myheap.Del(myh, 2)
	t.Log(myh)

	h := &IntHeap{}
	heap.Push(h, 10)
	heap.Push(h, 8)
	heap.Push(h, 11)
	heap.Push(h, 5)
	heap.Push(h, 4)
	heap.Push(h, 3)
	heap.Push(h, 20)
	t.Log(h)

	[]int(*h)[3] = 1
	heap.Fix(h, 3)
	[]int(*h)[2] = 12
	heap.Fix(h, 2)
	t.Log(h)

	heap.Remove(h, 2)
	t.Log(h)

	/* t.Log(heap.Pop(h))
	t.Log(heap.Pop(h))
	t.Log(heap.Pop(h))
	t.Log(h) */
}