题目描述

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

题目分析

二叉搜索树中 左 < 根 < 右，因此直接中序遍历，找到第 k 个被遍历到的数字即可。
可以递归，也可以用栈进行迭代遍历。

Java

private int count = 0;
private int ans = 0;

public int kthSmallest(TreeNode root, int k) {
    count = 0;
    dfs(root, k);
    return ans;
}

private void dfs(TreeNode root, int k) {
    if (root == null || count == k) {
        return;
    }
    dfs(root.left, k);
    count++;
    if (count == k) {
        ans = root.val;
    }
    dfs(root.right, k);
}

public int kthSmallest2(TreeNode root, int k) {
    Deque<TreeNode> stack = new LinkedList<>();
    TreeNode cur = root;
    int count = 0;
    while (cur != null || !stack.isEmpty()) {
        while (cur != null) {
            stack.push(cur);
            cur = cur.left;
        }
        cur = stack.pop();
        count++;
        if (count == k) {
            break;
        }
        cur = cur.right;
    }
    return cur.val;
}

Kotlin

private var count = 0
private var ans = 0

fun kthSmallest(root: TreeNode?, k: Int): Int {
    count = 0
    dfs(root, k)
    return ans
}

private fun dfs(root: TreeNode?, k: Int) {
    if (root == null || count == k) {
        return
    }
    dfs(root.left, k)
    count++
    if (count == k) {
        ans = root.`val`
    }
    dfs(root.right, k)
}

fun kthSmallest2(root: TreeNode?, k: Int): Int {
    val stack = LinkedList<TreeNode>()
    var cur = root
    var count = 0
    while (cur != null || !stack.isEmpty()) {
        while (cur != null) {
            stack.push(cur)
            cur = cur.left
        }
        cur = stack.pop()
        count++
        if (count == k) {
            return cur.`val`
        }
        cur = cur.right
    }
    return 0
}