advent-of-code-go/data-structures/heap/heap_test.go

package heap_test

import (
	"testing"

	"github.com/alexchao26/advent-of-code-go/data-structures/heap"
)

type mockNode int

func (n mockNode) Value() int {
	return int(n)
}

func TestMinHeap(t *testing.T) {
	h := heap.NewMinHeap()
	h.Add(mockNode(5))
	h.Add(mockNode(93))

	if h.Front().(mockNode).Value() != 5 {
		t.Errorf("After adding 5, h.Front().Value() = %d, want 5", h.Front().Value())
	}
	if h.Front().(mockNode).Value() != 5 {
		t.Errorf("After adding 5 & 93, h.Front() = %d, want 5", h.Front().Value())
	}

	// Add a bunch of nodes, make sure they are removed in order
	h.Add(mockNode(10))
	h.Add(mockNode(2))
	h.Add(mockNode(1))

	if h.Front().(mockNode).Value() != 1 {
		t.Errorf("After adding 5, 93, 10, 2 & 1, h.Front() = %d, want 1", h.Front().Value())
	}

	h.Add(mockNode(3))
	h.Add(mockNode(4))
	h.Add(mockNode(123))
	h.Add(mockNode(32))
	h.Add(mockNode(-15))

	// Ensure removing nodes returns in ascending order
	for _, want := range []int{-15, 1, 2, 3, 4, 5, 10, 32, 93, 123} {
		if got := h.Remove(); got.Value() != want {
			t.Errorf("h.Remove().Value() = %d, want %d", got.Value(), want)
		}
	}
}

func TestMaxHeap(t *testing.T) {
	h := heap.NewMaxHeap()
	h.Add(mockNode(5))
	h.Add(mockNode(93))

	if h.Front().Value() != 93 {
		t.Errorf("After adding 93, h.Front().Value() = %d, want 93", h.Front().Value())
	}
	if h.Front().Value() != 93 {
		t.Errorf("After adding 93 & 5, h.Front().Value() = %d, want 93", h.Front().Value())
	}

	// Add a bunch of nodes, make sure they are removed in order
	h.Add(mockNode(10))
	h.Add(mockNode(2))
	h.Add(mockNode(1))
	h.Add(mockNode(3))
	h.Add(mockNode(4))
	h.Add(mockNode(123))
	if h.Front().(mockNode).Value() != 123 {
		t.Errorf("After adding 5, 93, 10, 2, 1, 3, 4 & 123, h.Front() = %d, want 123", h.Front().Value())
	}

	h.Add(mockNode(32))
	h.Add(mockNode(-15))

	// Ensure removing returns in descending order
	for _, want := range []int{123, 93, 32, 10, 5, 4, 3, 2, 1, -15} {
		if got := h.Remove(); got.Value() != want {
			t.Errorf("h.Remove().Value() = %d, want %d", got.Value(), want)
		}
	}
}