题目描述

给定一个二叉搜索树的根节点 root ，和一个整数 k ，请你设计一个算法查找其中第 k 个最小元素（从 1 开始计数）。

示例 1：

输入： root = [3,1,4,null,2], k = 1
输出： 1

示例 2：

输入： root = [5,3,6,2,4,null,null,1], k = 3
输出： 3

提示：

• 树中的节点数为 n
• 1 <= k <= n <= 10⁴
• 0 <= Node.val <= 10⁴

代码及注释

/**
 * Definition for a binary tree node.
 * type TreeNode struct {
 *     Val int
 *     Left *TreeNode
 *     Right *TreeNode
 * }
 */

func kthSmallest(root *TreeNode, k int) int {
    // 初始化一个栈，用于迭代遍历二叉树
    var stack []*TreeNode
    // 迭代遍历二叉树
    for {
        // 将左子节点全部入栈
        for root != nil {
            stack = append(stack, root)
            root = root.Left
        }
        // 弹出栈顶元素
        root = stack[len(stack) - 1]
        stack = stack[:len(stack) - 1]
        // 找到第 k 小的元素
        k--
        if k == 0 {
            return root.Val
        }
        // 处理右子节点
        root = root.Right
    }
}

代码解释

• 使用中序遍历来遍历BST。中序遍历BST会按照递增顺序访问所有节点。
• 使用一个栈来迭代遍历二叉树。
• 当 k 减到 0 时，表示找到了第 k 小的元素，返回该元素的值。