Java——二叉搜索树中第k小的元素

题目链接

leetcode在线oj题——二叉搜索树中第k小的元素

题目描述

给定一个二叉搜索树的根节点 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^4
• 0 <= Node.val <= 10^4

解题思路

由于是二叉搜索树，所以我们可以直接中序遍历该树，然后从左到右找到第k个元素，即可得到第k小的元素

使用非递归的方式中序遍历：
定义一个栈stack，用于存储遍历过的节点，定义cur为当前遍历的节点

首先cur从根节点一直向左子树遍历，直到遍历到null，每遍历一个节点，将其压入stack

此时cur为null，从栈底到栈顶的元素从大到小排列，cur需要返回上一个遍历的位置，cur = stack弹出一个元素，k–，（如果k==0，直接返回当前节点值）

cur遍历完左子树后，开始遍历右子树，cur = cur.right，如果cur为null，让cur再次 = stack弹出一个元素

如果cur为空，需要返回上一次遍历的位置，如果所有位置都被遍历过了，没有返回的位置了，也就是栈为空了，循环结束
因此，循环的条件是：cur不等于空或者栈不为空

完整代码

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode() {}
 *     TreeNode(int val) { this.val = val; }
 *     TreeNode(int val, TreeNode left, TreeNode right) {
 *         this.val = val;
 *         this.left = left;
 *         this.right = right;
 *     }
 * }
 */
class Solution {
    public int kthSmallest(TreeNode root, int k) {
        Stack<TreeNode> stack = new Stack<>();
        TreeNode cur = root;
        
        do{
            while(cur != null) {
                stack.add(cur);
                cur = cur.left;
            }
            
            if(!stack.isEmpty()) {
                cur = stack.pop();
                k--;
                if(k == 0){
                    return cur.val;
                }
            }
            cur = cur.right; 
        }while(cur != null || !stack.isEmpty());
        
        return 0;
    }
}