invbintree/tree/tree.go

// Package tree implements a binary tree together with the inversion
// (mirror) operation and the traversals needed to observe it.
//
// Inversion swaps the left and right child of every node, producing the
// mirror image of the original tree. Two implementations are provided: a
// recursive one that is easiest to reason about, and an iterative one that
// runs in bounded stack space and so survives degenerate trees.
package tree

// Node is a single binary tree node. A nil *Node is a valid empty tree,
// which lets every operation below handle the empty case without a special
// branch at the call site.
type Node struct {
	Val         int
	Left, Right *Node
}

// New returns a leaf node holding val.
func New(val int) *Node { return &Node{Val: val} }

// Insert places val into the tree using binary-search-tree ordering and
// returns the (possibly new) root. Duplicate values go left, which keeps
// insertion total: no value is ever rejected.
func Insert(root *Node, val int) *Node {
	if root == nil {
		return New(val)
	}
	if val <= root.Val {
		root.Left = Insert(root.Left, val)
	} else {
		root.Right = Insert(root.Right, val)
	}
	return root
}

// FromSlice builds a tree by inserting vals in order.
func FromSlice(vals []int) *Node {
	var root *Node
	for _, v := range vals {
		root = Insert(root, v)
	}
	return root
}

// Invert mirrors the tree in place and returns the root.
//
// It recurses, so depth is bounded by the height of the tree. For trees that
// may be degenerate (a sorted insertion order produces a linked list of
// height n) prefer InvertIterative.
func Invert(root *Node) *Node {
	if root == nil {
		return nil
	}
	root.Left, root.Right = Invert(root.Right), Invert(root.Left)
	return root
}

// InvertIterative mirrors the tree in place using an explicit stack.
//
// This is the variant to reach for on untrusted input: a tree built from
// sorted data has height n, and the recursive version would need n stack
// frames. Here the heap-allocated stack grows instead, so a tall tree costs
// memory rather than a stack overflow.
func InvertIterative(root *Node) *Node {
	if root == nil {
		return nil
	}
	stack := []*Node{root}
	for len(stack) > 0 {
		n := stack[len(stack)-1]
		stack = stack[:len(stack)-1]
		n.Left, n.Right = n.Right, n.Left
		if n.Left != nil {
			stack = append(stack, n.Left)
		}
		if n.Right != nil {
			stack = append(stack, n.Right)
		}
	}
	return root
}

// Clone returns a deep copy, so callers can invert without mutating the
// original.
func Clone(root *Node) *Node {
	if root == nil {
		return nil
	}
	return &Node{Val: root.Val, Left: Clone(root.Left), Right: Clone(root.Right)}
}

// InOrder returns values left-root-right. For an unmodified BST this is
// sorted ascending; after inversion it is sorted descending, which is the
// cheapest observable proof that the mirror happened.
func InOrder(root *Node) []int {
	out := []int{}
	var walk func(*Node)
	walk = func(n *Node) {
		if n == nil {
			return
		}
		walk(n.Left)
		out = append(out, n.Val)
		walk(n.Right)
	}
	walk(root)
	return out
}

// PreOrder returns values root-left-right.
func PreOrder(root *Node) []int {
	out := []int{}
	var walk func(*Node)
	walk = func(n *Node) {
		if n == nil {
			return
		}
		out = append(out, n.Val)
		walk(n.Left)
		walk(n.Right)
	}
	walk(root)
	return out
}

// LevelOrder returns values breadth-first, one slice per depth level.
func LevelOrder(root *Node) [][]int {
	out := [][]int{}
	if root == nil {
		return out
	}
	level := []*Node{root}
	for len(level) > 0 {
		vals := make([]int, 0, len(level))
		var next []*Node
		for _, n := range level {
			vals = append(vals, n.Val)
			if n.Left != nil {
				next = append(next, n.Left)
			}
			if n.Right != nil {
				next = append(next, n.Right)
			}
		}
		out = append(out, vals)
		level = next
	}
	return out
}

// Height returns the number of nodes on the longest root-to-leaf path.
// An empty tree has height 0.
func Height(root *Node) int {
	if root == nil {
		return 0
	}
	l, r := Height(root.Left), Height(root.Right)
	if l > r {
		return l + 1
	}
	return r + 1
}

// Size returns the node count.
func Size(root *Node) int {
	if root == nil {
		return 0
	}
	return 1 + Size(root.Left) + Size(root.Right)
}

// Equal reports whether two trees have identical structure and values.
func Equal(a, b *Node) bool {
	if a == nil || b == nil {
		return a == b
	}
	return a.Val == b.Val && Equal(a.Left, b.Left) && Equal(a.Right, b.Right)
}

// IsMirror reports whether b is the mirror image of a. This is the
// definitional check for Invert, independent of the implementation.
func IsMirror(a, b *Node) bool {
	if a == nil || b == nil {
		return a == b
	}
	return a.Val == b.Val && IsMirror(a.Left, b.Right) && IsMirror(a.Right, b.Left)
}