leetcode/data-structure/ds-tree.go

package ds

// ds-tree.go
// 泛型二叉搜索树 (BST)。T 约束为 Ordered,Insert/Search 均基于比较。
// 中序遍历天然有序;提供递归与迭代两种遍历入口。
// 不做自平衡,退化时退化为链表;leetcode 题目里常需要这种朴素 BST 作为基础。

// TreeNode 二叉树节点。
type TreeNode[T any] struct {
	Val   T
	Left  *TreeNode[T]
	Right *TreeNode[T]
}

// BST 二叉搜索树。
type BST[T Ordered] struct {
	Root *TreeNode[T]
}

// Ordered 约束:支持 < <= > >= 比较。
type Ordered interface {
	~int | ~int8 | ~int16 | ~int32 | ~int64 |
		~uint | ~uint8 | ~uint16 | ~uint32 | ~uint64 | ~uintptr |
		~float32 | ~float64 | ~string
}

// NewBST 创建空 BST。
func NewBST[T Ordered]() *BST[T] { return &BST[T]{} }

// Insert 按 BST 性质插入,重复值放右子树。
func (t *BST[T]) Insert(v T) {
	t.Root = insertNode(t.Root, v)
}

func insertNode[T Ordered](node *TreeNode[T], v T) *TreeNode[T] {
	if node == nil {
		return &TreeNode[T]{Val: v}
	}
	if v < node.Val {
		node.Left = insertNode(node.Left, v)
	} else {
		node.Right = insertNode(node.Right, v)
	}
	return node
}

// Search 查找值是否存在。
func (t *BST[T]) Search(v T) bool {
	cur := t.Root
	for cur != nil {
		if v == cur.Val {
			return true
		}
		if v < cur.Val {
			cur = cur.Left
		} else {
			cur = cur.Right
		}
	}
	return false
}

// InOrder 中序遍历(左→根→右),BST 上结果升序。递归版。
func (t *BST[T]) InOrder() []T {
	var out []T
	var walk func(*TreeNode[T])
	walk = func(n *TreeNode[T]) {
		if n == nil {
			return
		}
		walk(n.Left)
		out = append(out, n.Val)
		walk(n.Right)
	}
	walk(t.Root)
	return out
}

// PreOrder 前序遍历(根→左→右)。递归版。
func (t *BST[T]) PreOrder() []T {
	var out []T
	var walk func(*TreeNode[T])
	walk = func(n *TreeNode[T]) {
		if n == nil {
			return
		}
		out = append(out, n.Val)
		walk(n.Left)
		walk(n.Right)
	}
	walk(t.Root)
	return out
}

// PostOrder 后序遍历(左→右→根)。递归版。
func (t *BST[T]) PostOrder() []T {
	var out []T
	var walk func(*TreeNode[T])
	walk = func(n *TreeNode[T]) {
		if n == nil {
			return
		}
		walk(n.Left)
		walk(n.Right)
		out = append(out, n.Val)
	}
	walk(t.Root)
	return out
}

// LevelOrder 层序遍历 (BFS),借助上面定义的 Queue。
func (t *BST[T]) LevelOrder() []T {
	if t.Root == nil {
		return nil
	}
	var out []T
	q := &Queue[*TreeNode[T]]{}
	q.Enqueue(t.Root)
	for !q.Empty() {
		node, _ := q.Dequeue()
		out = append(out, node.Val)
		if node.Left != nil {
			q.Enqueue(node.Left)
		}
		if node.Right != nil {
			q.Enqueue(node.Right)
		}
	}
	return out
}

// Height 树高(空树为 0,单节点为 1)。
func (t *BST[T]) Height() int {
	var h func(*TreeNode[T]) int
	h = func(n *TreeNode[T]) int {
		if n == nil {
			return 0
		}
		l := h(n.Left)
		r := h(n.Right)
		if l > r {
			return l + 1
		}
		return r + 1
	}
	return h(t.Root)
}

// DFS 迭代版深度优先遍历(前序),用栈模拟。作为迭代遍历的示例。
func (t *BST[T]) DFS() []T {
	if t.Root == nil {
		return nil
	}
	var out []T
	st := &Stack[*TreeNode[T]]{}
	st.Push(t.Root)
	for !st.Empty() {
		node, _ := st.Pop()
		out = append(out, node.Val)
		// 先压右再压左,弹出顺序才是 根→左→右
		if node.Right != nil {
			st.Push(node.Right)
		}
		if node.Left != nil {
			st.Push(node.Left)
		}
	}
	return out
}