go.charczuk.com@v0.0.0-20240327042549-bc490516bd1a/sdk/collections/binary_search_tree_test.go

/*

Copyright (c) 2023 - Present. Will Charczuk. All rights reserved.
Use of this source code is governed by a MIT license that can be found in the LICENSE file at the root of the repository.

*/

package collections

import (
	"cmp"
	"fmt"
	"strings"
	"testing"

	. "go.charczuk.com/sdk/assert"
)

func Test_BinarySearchTree_basic(t *testing.T) {
	bst := new(BinarySearchTree[int, string])
	for x := 0; x < 100; x++ {
		bst.Insert(makeKeyValue(x)())
	}

	minKey, minValue, ok := bst.Min()
	ItsEqual(t, true, ok)
	ItsEqual(t, 0, minKey)
	ItsEqual(t, "value-0", minValue)

	maxKey, maxValue, ok := bst.Max()
	ItsEqual(t, true, ok)
	ItsEqual(t, 99, maxKey)
	ItsEqual(t, "value-99", maxValue)

	for x := 0; x < 100; x++ {
		value, ok := bst.Search(x)
		ItsEqual(t, true, ok)
		ItsEqual(t, fmt.Sprintf("value-%d", x), value)
	}

	_, ok = bst.Search(-1)
	ItsEqual(t, false, ok)

	_, ok = bst.Search(101)
	ItsEqual(t, false, ok)

	for x := 0; x < 100; x++ {
		bst.Delete(x)
	}

	for x := 0; x < 100; x++ {
		_, ok = bst.Search(x)
		ItsEqual(t, false, ok)
	}

	// test insert overwrite
	bst.Insert(makeKeyValue(10)())
	bst.Insert(makeKeyValue(20)())

	bst.Insert(10, "not-value-10")

	value, ok := bst.Search(10)
	ItsEqual(t, true, ok)
	ItsEqual(t, "not-value-10", value)

	value, ok = bst.Search(20)
	ItsEqual(t, true, ok)
	ItsEqual(t, "value-20", value)
}

func Test_BinarySearchTree_acceptance(t *testing.T) {
	inputs := []int{
		81, 87, 47, 59, 83,
		18, 25, 40, 56, 0,
		94, 11, 62, 89, 28,
		74, 12, 45, 37, 6,
		95, 66, 29, 58, 101,
	}

	bst := new(BinarySearchTree[int, string])
	for _, x := range inputs {
		bst.Insert(makeKeyValue(x)())
	}

	node := func(k, height int, left, right *BinarySearchTreeNode[int, string]) *BinarySearchTreeNode[int, string] {
		return &BinarySearchTreeNode[int, string]{
			Key:    k,
			Height: height,
			Left:   left,
			Right:  right,
		}
	}

	goldenBST := &BinarySearchTree[int, string]{
		root: node(47, 5,
			node(25, 4,
				node(11, 3,
					node(0, 2, nil, node(6, 1, nil, nil)),
					node(18, 2, node(12, 1, nil, nil), nil),
				),
				node(40, 3,
					node(29, 2, node(28, 1, nil, nil), node(37, 1, nil, nil)),
					node(45, 1, nil, nil),
				),
			),
			node(81, 4,
				node(59, 3,
					node(56, 2, nil, node(58, 1, nil, nil)),
					node(66, 2,
						node(62, 1, nil, nil),
						node(74, 1, nil, nil),
					),
				),
				node(94, 3,
					node(87, 2,
						node(83, 1, nil, nil),
						node(89, 1, nil, nil),
					),
					node(95, 2, nil, node(101, 1, nil, nil)),
				),
			),
		),
	}
	ItsEqual(t, printTree(goldenBST), printTree(bst))
	ItsEqual(t, true, goldenBST.KeysEqual(bst), fmt.Sprintf("%s vs. %s", printTree(bst), printTree(goldenBST)))

	// delete 3 nodes, causing a rotation on the right side

	bst.Delete(94)
	bst.Delete(95)

	bst.Delete(89)
	bst.Delete(101)
	bst.Delete(83)

	goldenBST2 := &BinarySearchTree[int, string]{
		root: node(47, 5,
			node(25, 4,
				node(11, 3,
					node(0, 2, nil, node(6, 1, nil, nil)),
					node(18, 2, node(12, 1, nil, nil), nil),
				),
				node(40, 3,
					node(29, 2, node(28, 1, nil, nil), node(37, 1, nil, nil)),
					node(45, 1, nil, nil),
				),
			),
			node(59, 4,
				node(56, 2,
					nil,
					node(58, 1, nil, nil),
				),
				node(81, 3,
					node(66, 2,
						node(62, 1, nil, nil),
						node(74, 1, nil, nil),
					),
					node(87, 1, nil, nil),
				),
			),
		),
	}

	ItsEqual(t, printTree(goldenBST2), printTree(bst))
	ItsEqual(t, true, goldenBST2.KeysEqual(bst), fmt.Sprintf("%s vs. %s", printTree(bst), printTree(goldenBST2)))
}

func Test_BinarySearchTree_InOrder(t *testing.T) {
	bst := new(BinarySearchTree[int, string])
	for x := 0; x < 5; x++ {
		bst.Insert(makeKeyValue(x)())
	}
	ItsEqual(t, 1, bst.root.Key)

	var keys []int
	bst.InOrder(func(k int, _ string) {
		keys = append(keys, k)
	})
	ItsEqual(t, []int{0, 1, 2, 3, 4}, keys)
}

func Test_BinarySearchTree_PreOrder(t *testing.T) {
	bst := new(BinarySearchTree[int, string])
	for x := 0; x < 5; x++ {
		bst.Insert(makeKeyValue(x)())
	}
	ItsEqual(t, 1, bst.root.Key)

	var keys []int
	bst.PreOrder(func(k int, _ string) {
		keys = append(keys, k)
	})
	ItsEqual(t, []int{1, 0, 3, 2, 4}, keys)
}

func Test_BinarySearchTree_PostOrder(t *testing.T) {
	bst := new(BinarySearchTree[int, string])
	for x := 0; x < 5; x++ {
		bst.Insert(makeKeyValue(x)())
	}
	ItsEqual(t, 1, bst.root.Key)

	var keys []int
	bst.PostOrder(func(k int, _ string) {
		keys = append(keys, k)
	})
	ItsEqual(t, []int{0, 2, 4, 3, 1}, keys)
}

func Test_BinarySearchTree_unbalancedLeft(t *testing.T) {
	bst := new(BinarySearchTree[int, string])
	for x := 0; x < 100; x++ {
		bst.Insert(makeKeyValue(x)())
	}

	for x := 0; x < 50; x++ {
		bst.Delete(x)
	}

	ItsEqual(t, 79, bst.root.Key)
}

func Test_BinarySearchTree_reverse(t *testing.T) {
	bst := new(BinarySearchTree[int, string])
	for x := 99; x >= 0; x-- {
		bst.Insert(makeKeyValue(x)())
	}
	ItsEqual(t, nil, verify(bst.root))
	ItsEqual(t, 36, bst.root.Key)
}

//
// bst test helpers
//

func makeKeyValue(x int) func() (int, string) {
	return func() (int, string) {
		return x, fmt.Sprintf("value-%d", x)
	}
}

func printTree[K cmp.Ordered, V any](bst *BinarySearchTree[K, V]) string {
	var keys []string
	bst.InOrder(func(k K, _ V) {
		keys = append(keys, fmt.Sprint(k))
	})
	return "[ " + strings.Join(keys, " ") + " ]"
}

func verify[K cmp.Ordered, V any](n *BinarySearchTreeNode[K, V]) error {
	if n.Left != nil {
		min := verifyMin(n)
		if min.Key >= n.Key {
			return fmt.Errorf("invalid min; %v is greater than parent %v", min.Key, n.Key)
		}
		if err := verify(n.Left); err != nil {
			return err
		}
	}
	if n.Right != nil {
		max := verifyMax(n)
		if max.Key <= n.Key {
			return fmt.Errorf("invalid max; %v is less than parent %v", max.Key, n.Key)
		}
		if err := verify(n.Right); err != nil {
			return err
		}
	}
	return nil
}

func verifyMin[K cmp.Ordered, V any](n *BinarySearchTreeNode[K, V]) (min *BinarySearchTreeNode[K, V]) {
	min = n
	for min.Left != nil {
		min = min.Left
	}
	return
}

func verifyMax[K cmp.Ordered, V any](n *BinarySearchTreeNode[K, V]) (max *BinarySearchTreeNode[K, V]) {
	max = n
	for max.Right != nil {
		max = max.Right
	}
	return
}

func Test_max(t *testing.T) {
	v := max[int]()
	ItsEqual(t, 0, v)

	v = max(1)
	ItsEqual(t, 1, v)

	v = max(1, 2, 3, 4)
	ItsEqual(t, 4, v)
}

func Test_keysEqual(t *testing.T) {
	bst := new(BinarySearchTree[int, func()])
	ItsEqual(t, true, bst._keysEqual(nil, nil))

	a := &BinarySearchTreeNode[int, func()]{
		Key: 1,
	}

	ItsEqual(t, false, bst._keysEqual(a, nil))
	ItsEqual(t, false, bst._keysEqual(nil, a))

	b := &BinarySearchTreeNode[int, func()]{
		Key: 2,
	}
	ItsEqual(t, false, bst._keysEqual(a, b))

	a1 := &BinarySearchTreeNode[int, func()]{
		Key:    1,
		Height: 10,
	}
	ItsEqual(t, false, bst._keysEqual(a, a1))

	a0 := &BinarySearchTreeNode[int, func()]{
		Key: 1,
	}
	ItsEqual(t, true, bst._keysEqual(a, a0))
}