LeetCode 第 230 题：“二叉搜索树中第 K 小的元素”题解
题解地址：递归与非递归写法（同理完成第 144、94、145 题，Python 代码、Java 代码）。
说明：文本首发在力扣的题解版块，更新也会在第 1 时间在上面的网站中更新，这篇文章只是上面的文章的一个快照，您可以点击上面的链接看到其他网友对本文的评论。
传送门：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 函数？
递归与非递归写法（同理完成第 144、94、145 题，Python 代码、Java 代码）
思路分析：利用“二叉搜索树”在“中序遍历”以后，得到的是有序数组，那么我们就中序遍历好了，遍历到第 个数即可。
我写下来发现递归的写法比较容易写错，要设置全局变量，而非递归的写法还相对比较“通用”且好理解。
参考代码 1：使用递归“中序遍历”。
Python 代码：
class TreeNode:
def __init__(self, x):
self.val = x
self.left = None
self.right = None
# 使用递归的方法，中序遍历
class Solution:
def __init__(self):
self.counter = 0
self.res = 0
def kthSmallest(self, root, k):
# 递归执行左子树的逻辑
if root.left:
# 不是空，才继续遍历
self.kthSmallest(root.left, k)
# 在这里执行操作，数到第 k 个即可
self.counter += 1
# print(root.val)
if self.counter == k:
# 注意：千万不能在这里返回，后序遍历还要继续进行下去
self.res = root.val
# 注意：这里不能加 return
# 递归执行右子树的逻辑
if root.right:
self.kthSmallest(root.right, k)
return self.res
if __name__ == '__main__':
node3 = TreeNode(3)
node1 = TreeNode(1)
node4 = TreeNode(4)
node2 = TreeNode(2)
node3.left = node1
node3.right = node4
node1.right = node2
solution = Solution()
result = solution.kthSmallest(node3, k=1)
print(result)
Java 代码：
class TreeNode {
int val;
TreeNode left;
TreeNode right;
TreeNode(int x) {
val = x;
}
}
// 解题关键：中序遍历
// https://leetcode-cn.com/problems/kth-smallest-element-in-a-bst/description/
// 给定一个二叉搜索树，编写一个函数kthSmallest来查找其中第 k 个最小的元素。
// 只要利用二分搜索树的中序遍历，就可以完成。
public class Solution {
private int count = 0;
private int res = 0;
private void dfs(TreeNode node) {
if (node == null) {
// 什么都不做
return;
}
dfs(node.left);
count--;
if (count == 0) {
this.res = node.val;
}
dfs(node.right);
}
// k 如果在方法传递的过程中是值传递，所以把它设置为成员变量，这样就是引用传递
// 因为我们要用到 k 全局的值，去数出，我是第几个中序遍历到的值
public int kthSmallest(TreeNode root, int k) {
count = k;
dfs(root);
return res;
}
public static void main(String[] args) {
TreeNode treeNode1 = new TreeNode(10);
TreeNode treeNode2 = new TreeNode(15);
TreeNode treeNode3 = new TreeNode(20);
treeNode2.left = treeNode1;
treeNode2.right = treeNode3;
Solution solution = new Solution();
int kthSmallest = solution.kthSmallest(treeNode2, 2);
System.out.println(kthSmallest);
}
}
下面是 Python 的另一种写法：使用 global 关键字，仍需使用辅助函数：
Python 代码：
class TreeNode:
def __init__(self, x):
self.val = x
self.left = None
self.right = None
class Solution:
def kthSmallest(self, root, k):
global counter, res
counter = 0
res = 0
def dfs(root, k):
if not root:
# 如果是空，直接退出
return
dfs(root.left, k)
global counter, res
counter += 1
if counter == k:
res = root.val
dfs(root.right, k)
dfs(root, k)
return res
参考代码 2：模拟系统栈的方式：使用二叉树非递归遍历的通用方法。使用同样的的方法还可以解决 「力扣」LeetCode 第 144 题：二叉树的前序遍历、「力扣」第 94 题：二叉树的中序遍历、「力扣」第 145 题：二叉树的后序遍历。
Python 代码：
class TreeNode:
def __init__(self, x):
self.val = x
self.left = None
self.right = None
class Solution:
# 模拟系统栈的方式实现，是一种比较通用的做法，
# 可以作为二叉树的三种非递归遍历
def kthSmallest(self, root, k):
# 0 表示当前遍历到它，1 表示压入栈
# 刚开始是 1 ，不要写成 0 了
stack = [(1, root)]
while stack:
command, node = stack.pop()
if node is None:
# 不能写 return ，这不是递归
continue
if command == 0:
k -= 1
if k == 0:
return node.val
else:
# 此时 command == 1 的时候，表示递归遍历到的
# 注意：写的时候倒过来写
stack.append((1, node.right))
stack.append((0, node))
stack.append((1, node.left))
其实入栈的时候，就可以判断，我们只将非空结点入栈，推荐下面这种写法：
Python 代码：
class Solution:
def kthSmallest(self, root, k):
stack = [(1, root)]
while stack:
command, node = stack.pop()
if command == 0:
k -= 1
if k == 0:
return node.val
else:
# 模拟系统栈实现中序遍历(先左边、再自己、再右边)
# 注意：写的时候倒过来写
if node.right:
stack.append((1, node.right))
stack.append((0, node))
if node.left:
stack.append((1, node.left))
Java 代码：
import java.util.Stack;
public class Solution4 {
private enum Action {
// GO 表示递归处理
// ADDTORESULT 表示当前马上执行将结点的值添加到结果集中
GO, ADDTORESULT
}
private class Command {
private Action action;
private TreeNode node;
public Command(Action action, TreeNode node) {
this.action = action;
this.node = node;
}
}
public int kthSmallest(TreeNode root, int k) {
Stack<Command> stack = new Stack<>();
stack.add(new Command(Action.GO, root));
while (!stack.isEmpty()) {
Command cur = stack.pop();
TreeNode node = cur.node;
if (cur.action == Action.ADDTORESULT) {
k--;
if (k == 0) {
return node.val;
}
} else {
assert cur.action == Action.GO;
if (node.right != null) {
stack.add(new Command(Action.GO, node.right));
}
stack.add(new Command(Action.ADDTORESULT, node));
if (node.left != null) {
stack.add(new Command(Action.GO, node.left));
}
}
}
throw new RuntimeException("参数错误");
}
}