题目描述

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

(DFS中序遍历) $O(h + k)$

因为二叉搜索树的中序遍历序列就是有序的，而第k小的元素就是中序遍历到的第k个元素，所以我们可以用递归来中序遍历这棵树，记录当前已经遍历过的节点数量，同时用一个全局变量来记录答案。当我们已经找到答案时，我们无需遍历该节点的右子树，可以直接返回答案。

时间复杂度

在中序遍历时我们会先沿根节点的左子树一直向左来访问第一个节点，这需要 $O(h)$ 的时间，这里 $h$ 是树的高度，之后需要 $O(k)$ 的时间来找到第k个节点。

C++ 代码

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    TreeNode* ans;
    int kthSmallest(TreeNode* root, int k) {
        ans = nullptr;
        inorder(root, k);
        return ans->val;
    }

    void inorder(TreeNode* root, int& k) {
        if (!root) return;
        inorder(root->left, k);
        k -- ;
        if (!k) ans = root;
        if (k > 0) inorder(root->right, k);
    }
};

算法2

(栈+中序遍历) $O(h + k)$

当然我们也可以用栈来模拟中序遍历的过程，当我们访问到第k个节点时，我们就找到了答案。

时间复杂度

时间复杂度与DFS相同，不过当树的层数很深时用栈不会像DFS一样造成调用栈爆栈。

C++代码

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    int kthSmallest(TreeNode* root, int k) {
        stack<TreeNode*> stk;

        while (stk.size() || root) {
            while (root) {
                stk.push(root);
                root = root->left;
            }

            root = stk.top();
            stk.pop();
            if (-- k == 0) return root->val;
            root = root->right;
        }
        return 0;
    }
};

进阶

如果该树经常被修改，那么树可能会变成不平衡的，即树高由 $O(\log n)$ 变为 $O(n)$。使用Splay树或者Treap等平衡二叉树可以保证树高 $h$ 是 $O(\log n)$级别的，如果我们在结点中增加一个cnt属性来记录以该节点为根的子树的节点个数，那么我们就可以在 $O(h)$ 的时间内找到第k大的元素。