Leetcode: 230. 二叉搜索树中第K小的元素

题目描述

一个二叉搜索树，编写一个函数 kthSmallest 来查找其中第 k 个最小的元素。

说明：

你可以假设 k 总是有效的，1 ≤ k ≤ 二叉搜索树元素个数。

示例 1:

输入: root = [3,1,4,null,2], k = 1
3
/ \
1   4
\
   2
输出: 1

示例 2:

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

进阶：

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

思路

思路1 中序遍历+剪枝+递归

直接递归使用中序遍历，然后通过设置一个计数器计算数在整个二叉搜索中排序的位置，到达第K小时返回结果（超过K直接返回）

思路2 中序遍历+剪枝+迭代

使用一个栈实现中序遍历迭代，通过出栈的数量找到第K个数

代码

代码1

class Solution {
    int numCount;
    int result;
    
    public int kthSmallest(TreeNode root, int k) {
        numCount = 0;
        addSearchNum(root, k);
        return result;
    }
    
    private void addSearchNum(TreeNode root, int k){
        if (root == null || numCount > k){
            return;
        }
        addSearchNum(root.left, k);
        numCount ++;
        if (numCount == k){
            result = root.val;
        }
        addSearchNum(root.right, k);
    }
}

代码2

class Solution {
    public int kthSmallest(TreeNode root, int k) {
        Stack stack = new Stack<>();
        while(true){
            while (root != null){
                stack.push(root);
                root = root.left;
            }
            root = stack.pop();
            k --;
            if (k == 0) return root.val;
            root = root.right;
        }
    }
}