009.二叉搜索树中第K小的元素

• 230.二叉搜索树中第K小的元素
  • 中序遍历
  • 记录子树的结点数
  • 平衡二叉搜索树

相关题目：

• 1382.将二叉搜索树变平衡
  • 中序遍历后构造
  • 手撸 AVL
    • 通用性更强，不要求是二叉搜索树

二叉搜索树中第 K 小的元素 - 记录子树的结点数

如果你需要频繁地查找第 k 小的值，你将如何优化算法？

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
 
var counts map[*TreeNode]int
 
func kthSmallest(root *TreeNode, k int) int {
	counts = map[*TreeNode]int{}
	buildCounts(root)
	node := searchKthSmallest(root, k)
	if root == nil {
		panic("k too big")
	}
	return node.Val
}
 
func buildCounts(root *TreeNode) int {
	if root == nil {
		return 0
	}
	counts[root] = 1 + buildCounts(root.Left) + buildCounts(root.Right)
	return counts[root]
}
 
func searchKthSmallest(root *TreeNode, k int) *TreeNode {
	for {
		if root == nil || counts[root] < k {
			return nil
		}
		leftNodeCount := counts[root.Left]
		if leftNodeCount == k - 1 {
			return root
		}
		if leftNodeCount < k - 1 {
			root = root.Right
			k -= leftNodeCount + 1
		} else {
			root = root.Left
		}
	}
}

二叉搜索树中第 K 小的元素 - 平衡二叉搜索树

如果二叉搜索树经常被修改（插入/删除操作）并且你需要频繁地查找第 k 小的值，你将如何优化算法？

我们注意到在 记录子树的结点数方法 中搜索二叉搜索树的时间复杂度为 O(H)，其中 H 是树的高度；当树是平衡树时，时间复杂度取得最小值 O(log⁡N)。因此，我们在记录子树的结点数的基础上，将二叉搜索树转换为平衡二叉搜索树，并在插入和删除操作中维护它的平衡状态。

其中，将二叉搜索树转换为平衡二叉搜索树，可以参考 将二叉搜索树变平衡。在插入和删除操作中维护平衡状态相对复杂，即： todo 手撸 AVL。

将二叉搜索树变平衡 - 中序遍历后构造

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */
func balanceBST(root *TreeNode) *TreeNode {
	return buildAVL(inorder(root))
}
 
func inorder(root *TreeNode) []*TreeNode {
	res := []*TreeNode{}
	stack := []*TreeNode{}
	cur := root
 
	for cur != nil || len(stack) > 0 {
		if cur != nil {
			stack = append(stack, cur)
			cur = cur.Left
		} else {
			cur = stack[len(stack) - 1]
			stack = stack[:len(stack) - 1]
			res = append(res, cur)
			cur = cur.Right
		}
	}
 
	return res
}
 
func buildAVL(arr []*TreeNode) *TreeNode {
	n := len(arr)
	if n == 0 {
		return nil
	}
	node := arr[n >> 1]
	node.Left = buildAVL(arr[:n >> 1])
	node.Right = buildAVL(arr[n >> 1 + 1:])
	return node
}