dataStructures/trees/avl/AvlTree.go

package avl

import (
	"cmp"
	"fmt"

	"datastructures/linear"
)

type AvlNode[T cmp.Ordered] struct {
	Left   *AvlNode[T]
	Right  *AvlNode[T]
	Data   T
	Height int
}

type AvlTree[T cmp.Ordered] struct {
	Root *AvlNode[T]
}

// NewAvlTree() -> creates an AVL tree
func NewAvlTree[T cmp.Ordered]() *AvlTree[T] {
	return &AvlTree[T]{}
}

// height() ->  0 if root
func height[T cmp.Ordered](n *AvlNode[T]) int {
	if n == nil {
		return 0
	}
	return n.Height
}

// updateHeight()
func updateHeight[T cmp.Ordered](node *AvlNode[T]) {
	node.Height = 1 + max(height(node.Left), height(node.Right))
}

// balanceFactor() -> only 0 , -1 ,-2 ok
func balanceFactor[T cmp.Ordered](n *AvlNode[T]) int {
	if n == nil {
		return 0
	}

	return height(n.Left) - height(n.Right)
}

// rotateRight() -> fixes a left heavy subtree, returns the new subtree root
func rotateRight[T cmp.Ordered](y *AvlNode[T]) *AvlNode[T] {
	x := y.Left
	b := x.Right

	x.Right = y
	y.Left = b

	updateHeight(y)
	updateHeight(x)

	return x
}

// rotateLeft() -> fixes a right heavy subtree, returns the new subtree root
func rotateLeft[T cmp.Ordered](x *AvlNode[T]) *AvlNode[T] {
	y := x.Right
	b := y.Left

	y.Left = x
	x.Right = b

	updateHeight(x)
	updateHeight(y)

	return y
}

// Insert() -> inserts data, keeping the tree balanced
func (tree *AvlTree[T]) Insert(data T) {
	tree.Root = insertHelper(tree.Root, data)
}

func insertHelper[T cmp.Ordered](node *AvlNode[T], data T) *AvlNode[T] {
	//  normal BST insert
	if node == nil {
		return &AvlNode[T]{Data: data, Height: 1}
	}

	if data < node.Data {
		node.Left = insertHelper(node.Left, data)
	} else if data > node.Data {
		node.Right = insertHelper(node.Right, data)
	} else {
		return node // duplicate, ignore
	}

	updateHeight(node)

	bf := balanceFactor(node)

	if bf > 1 {
		if data > node.Left.Data { // Left-Right case
			node.Left = rotateLeft(node.Left)
		}
		return rotateRight(node) // Left-Left case
	}

	if bf < -1 {
		if data < node.Right.Data {
			node.Right = rotateRight(node.Right)
		}
		return rotateLeft(node)
	}

	return node
}

// Find() => binary search implementation returns a pointer to the node and true
// if val is found in tree, else returns zero struct and false
func (t AvlTree[T]) Find(val T) (*AvlNode[T], bool) {
	zero := &AvlNode[T]{}

	if t.Root == nil {
		return zero, false
	}

	for t.Root != nil {
		if val == t.Root.Data {
			return t.Root, true
		} else if val < t.Root.Data {
			t.Root = t.Root.Left
		} else if val > t.Root.Data {
			t.Root = t.Root.Right
		}
	}

	return zero, false
}

// Traversals

// TraverseInorder() -> in order traversal of avl tree starting at root, returns a slice of
// AvlNode.Data. Iterative, no recursion.
func (tree *AvlTree[T]) TraverseInorder() []T {
	var list []T
	stack := linear.Stack[*AvlNode[T]]{}

	current := tree.Root

	// every node on the stack has been walked past but not emitted yet
	for current != nil || stack.Size() > 0 {

		// dive left
		for current != nil {
			stack.Push(current)
			current = current.Left
		}

		node, err := stack.Pop()
		if err != nil {
			break
		}

		list = append(list, node.Data)

		current = node.Right
	}

	return list
}

// TraversePreorder() -> pre order traversal of avl tree starting at root, returns a slice of
// AvlNode.Data
func (tree *AvlTree[T]) TraversePreorder() []T {
	var list []T
	stack := linear.Stack[*AvlNode[T]]{}

	if tree.Root == nil {
		return list
	}

	stack.Push(tree.Root)

	for stack.Size() > 0 {

		root, err := stack.Pop()
		if err != nil {
			break
		}

		list = append(list, root.Data)

		if root.Right != nil {
			stack.Push(root.Right)
		}

		if root.Left != nil {
			stack.Push(root.Left)
		}

	}

	return list
}

// Display() -> draws the tree rotated 90° left: right subtree on top,
// left subtree on the bottom, connected with box-drawing characters.
// Iterative , no recursion.
func (t AvlTree[T]) Display() {
	if t.Root == nil {
		fmt.Println("<empty>")
		return
	}

	// A frame is either a node to expand or a ready made line to print.
	type frame struct {
		node   *AvlNode[T]
		prefix string
		isLeft bool
		isRoot bool
		line   string
		print  bool
	}

	stack := []frame{{node: t.Root, isRoot: true}}

	for len(stack) > 0 {
		f := stack[len(stack)-1]
		stack = stack[:len(stack)-1]

		if f.print {
			fmt.Println(f.line)
			continue
		}

		n := f.node

		var connector, rightPrefix, leftPrefix string
		switch {
		case f.isRoot:
			connector = ""
			rightPrefix, leftPrefix = "    ", "    "
		case f.isLeft:
			connector = "└── "
			rightPrefix, leftPrefix = f.prefix+"│   ", f.prefix+"    "
		default:
			connector = "┌── "
			rightPrefix, leftPrefix = f.prefix+"    ", f.prefix+"│   "
		}

		line := f.prefix + connector + fmt.Sprintf("%v", n.Data)

		if n.Left != nil {
			stack = append(stack, frame{node: n.Left, prefix: leftPrefix, isLeft: true})
		}
		stack = append(stack, frame{print: true, line: line})
		if n.Right != nil {
			stack = append(stack, frame{node: n.Right, prefix: rightPrefix, isLeft: false})
		}
	}
}