leetcode [230] 二叉搜索树中第K小的元素

---

Contact me:

Blog -> https://cugtyt.github.io/blog/index
Email -> cugtyt@qq.com
GitHub -> Cugtyt@GitHub

---

给定一个二叉搜索树，编写一个函数 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 函数？

由于是搜索树，递归遍历就是顺序，累计第k个的时候提前终止。

class Solution:
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        ans = 0
        n = 0
        def dfs(root):
            nonlocal n, ans
            if not root or n >= k:
                return
            dfs(root.left)
            n += 1
            if n == k:
                ans = root.val
                return
            dfs(root.right) 
        dfs(root)
        return ans