heapq.go

heapq.go

package heapq

import "errors"

type intMaxHeap struct {
	buf    []int
	length int
}

func (h *intMaxHeap) Cap() int {
	return cap(h.buf)
}

func (h *intMaxHeap) Len() int {
	return h.length
}

func (h *intMaxHeap) Push(item int) error {
	if h.length == cap(h.buf) {
		return errors.New("heap is full")
	}

	n := h.length
	h.buf[n] = item
	h.length++

	for n != 0 {
		m := (n - 1) / 2
		parent := h.buf[m]

		if item <= parent {
			return nil
		}

		h.buf[m] = item
		h.buf[n] = parent

		n = m
	}
	return nil
}

func (h *intMaxHeap) Pop() (int, error) {
	if h.length == 0 {
		return 0, errors.New("heap is empty")
	}

	largest := h.buf[0]

	h.buf[0] = h.buf[h.length-1]
	h.length--

	n := 0

	for {
		l := 2*(n+1) - 1
		r := 2 * (n + 1)

		m := r
		if r >= h.length {
			if l >= h.length {
				return largest, nil
			}
			m = l
		}
		if h.buf[l] > h.buf[r] {
			m = l
		}

		if h.buf[n] >= h.buf[m] {
			return largest, nil
		}
		tmp := h.buf[n]
		h.buf[n] = h.buf[m]
		h.buf[m] = tmp

		n = m
	}
	panic("should not reach here")
	return largest, nil
}

func (h *intMaxHeap) Largest() int {
	return h.buf[0]
}

func NewIntMaxHeap(c int) *intMaxHeap {
	buf := make([]int, c)
	return &intMaxHeap{buf: buf, length: 0}
}
    

heapq_test.go

package heapq

import (
	"math/rand"
	"testing"
)

func TestNewIntMaxHeap(t *testing.T) {
	h := NewIntMaxHeap(10)

	actual := h.Len()
	if actual != 0 {
		t.Errorf("length of a new heap must be zero, got %d", actual)
	}

	expected := 10
	actual = h.Cap()

	if actual != expected {
		t.Errorf("capacity of a new heap must be that was initialized", actual, expected)
	}
}

func TestPush(t *testing.T) {
	h := NewIntMaxHeap(10)

	h.Push(4)
	h.Push(6)
	h.Push(5)

	expected := 6
	actual := h.Largest()

	if actual != expected {
		t.Errorf("wrong largest value, got %v, expect %v", actual, expected)
	}
}

func TestPop(t *testing.T) {
	h := NewIntMaxHeap(10)

	h.Push(4)
	h.Push(6)
	h.Push(5)
	expected := 6
	actual, err := h.Pop()

	if err != nil {
		t.Errorf("error in Pop: %s", err.Error())
	}

	if actual != expected {
		t.Errorf("not the largest value, got %v, expect %v", actual, expected)
	}

	expected = 5
	actual = h.Largest()

	if actual != expected {
		t.Errorf("not the largest value, got %v, expect %v", actual, expected)
	}
}

func BenchmarkTop10(b *testing.B) {
	N := 10
	h := NewIntMaxHeap(N)
	b.ResetTimer()
	benchBody(b, N, h)
}

func BenchmarkTop100(b *testing.B) {
	N := 100
	h := NewIntMaxHeap(N)
	b.ResetTimer()
	benchBody(b, N, h)
}

func BenchmarkTop1000(b *testing.B) {
	N := 1000
	h := NewIntMaxHeap(N)
	b.ResetTimer()
	benchBody(b, N, h)
}

func BenchmarkTop10000(b *testing.B) {
	N := 10000
	h := NewIntMaxHeap(N)
	b.ResetTimer()
	benchBody(b, N, h)
}

func BenchmarkTop100000(b *testing.B) {
	N := 100000
	h := NewIntMaxHeap(N)
	b.ResetTimer()
	benchBody(b, N, h)
}

func BenchmarkTop1000000(b *testing.B) {
	N := 1000000
	h := NewIntMaxHeap(N)
	b.ResetTimer()
	benchBody(b, N, h)
}

func BenchmarkTop10000000(b *testing.B) {
	N := 10000000
	h := NewIntMaxHeap(N)
	b.ResetTimer()
	benchBody(b, N, h)
}

func benchBody(b *testing.B, N int, h *intMaxHeap) {
	for _, n := range rand.Perm(N) {
		h.Push(n)
	}

	for i := 0; i < 10; i++ {
		actual, err := h.Pop()
		if err != nil {
			b.Errorf("error in Pop: %s", err.Error())
		}
		expected := N - i - 1
		if actual != expected {
			b.Errorf("not the largest value, got %v, expect %v", actual, expected)
		}
	}
}