[每日一题] 134. 二叉搜索树中第K小的元素(BST树、递归、搜索技巧、多方法）

1. 题目来源

链接：二叉搜索树中第K小的元素
来源：LeetCode

2. 题目说明

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

3. 题目解析

方法一：BST性质、非递归中序遍历

用 BST 的性质来解题，最重要的性质是就是左<根<右，如果用中序遍历所有的节点就会得到一个有序数组。所以解题的关键还是中序遍历啊。

先来看一种非递归的方法，中序遍历最先遍历到的是最小的结点，只要用一个计数器，每遍历一个结点，计数器自增1，当计数器到达 k 时，返回当前结点值即可。

参见代码如下：

// 执行用时 :24 ms, 在所有 C++ 提交中击败了70.31%的用户
// 内存消耗 :21.6 MB, 在所有 C++ 提交中击败了33.58%的用户

class Solution {
public:
    int kthSmallest(TreeNode* root, int k) {
        int cnt = 0;
        stack<TreeNode*> s;
        TreeNode *p = root;
        while (p || !s.empty()) {
            while (p) {
                s.push(p);
                p = p->left;
            }
            p = s.top(); 
            s.pop();
            ++cnt;
            if (cnt == k) 
                return p->val;
            p = p->right;
        }
        return 0;
    }
};

方法二：递归中序遍历

当然，此题也可以用递归来解，还是利用中序遍历来解。

参见代码如下：

// 执行用时 :28 ms, 在所有 C++ 提交中击败了44.03%的用户
// 内存消耗 :22 MB, 在所有 C++ 提交中击败了6.89%的用户

/**
 * 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 InOrder(TreeNode* root, int& k) {
        if (root) {
            int val = InOrder(root->left, k);
            if (k == 0) 
                return val;
            if (--k == 0) 
                return root->val;
            return InOrder(root->right, k);
        }
        return -1;
    }

    int kthSmallest(TreeNode* root, int k) {
        return InOrder(root, k);
    }
};

方法三：分治法

来看一种分治法的思路，由于 BST 的性质，可以快速定位出第 k 小的元素是在左子树还是右子树，即有：

• 首先计算出左子树的结点个数总和 cnt
• 如果 k 小于等于左子树结点总和 cnt，说明第 k 小的元素在左子树中，直接对左子结点调用递归即可
• 如果 k 大于 cnt + 1，说明目标值在右子树中，对右子结点调用递归函数

注意此时的 k 应为 k - cnt - 1，因为已经减少了 cnt + 1 个结点。如果 k 正好等于 cnt + 1，说明当前结点即为所求，返回当前结点值即可。

参见代码如下：

// 执行用时 :20 ms, 在所有 C++ 提交中击败了88.64%的用户
// 内存消耗 :22.1 MB, 在所有 C++ 提交中击败了5.23%的用户

/**
 * 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) {
        int cnt = count(root->left);
        if (k <= cnt) 
            return kthSmallest(root->left, k);
        else if (k > cnt + 1) 
            return kthSmallest(root->right, k - cnt - 1);
        return root->val;
    }
    int count(TreeNode* node) {
        if (!node) 
            return 0;
        return 1 + count(node->left) + count(node->right);
    }
};

方法四：统计左右子树节点个数、探索二叉搜索树解法（绝妙）

这道题的 进阶 中说假设该 BST 被修改的很频繁，而且查找第k小元素的操作也很频繁，问我们如何优化。

其实最好的方法还是像上面的解法那样利用分治法来快速定位目标所在的位置，但是每个递归都遍历左子树所有结点来计算个数的操作并不高效，所以应该修改原树结点的结构，使其保存包括当前结点和其左右子树所有结点的个数，这样就可以快速得到任何左子树结点总数来快速定位目标值了。

定义了新结点结构体，然后就要生成新树，还是用递归的方法生成新树，注意生成的结点的 count 值要累加其左右子结点的 count 值。

然后在求第 k 小元素的函数中，先生成新的树，然后调用递归函数。在递归函数中，不能直接访问左子结点的 count 值，因为左子节结点不一定存在，所以要先判断，如果左子结点存在的话，那么跟上面解法的操作相同。如果不存在的话，当此时 k 为 1 的时候，直接返回当前结点值，否则就对右子结点调用递归函数，k自减1。

在 LeetCode探索 二叉搜索树 对该种搜索方法进行了诠释，真的绝妙的思想。

参见代码如下：

// 执行用时 :32 ms, 在所有 C++ 提交中击败了24.55%的用户
// 内存消耗 :24.9 MB, 在所有 C++ 提交中击败了5.23%的用户

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */


/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */


 // Follow up
class Solution {
public:
    struct MyTreeNode {
        int val;
        int count;
        MyTreeNode *left;
        MyTreeNode *right;
        MyTreeNode(int x) : val(x), count(1), left(NULL), right(NULL) {}
    };
    
    MyTreeNode* build(TreeNode* root) {
        if (!root) 
            return NULL;
        MyTreeNode* node = new MyTreeNode(root->val);
        node->left = build(root->left);
        node->right = build(root->right);
        if (node->left) 
            node->count += node->left->count;
        if (node->right) 
            node->count += node->right->count;
        return node;
    }
    
    int kthSmallest(TreeNode* root, int k) {
        MyTreeNode *node = build(root);
        return helper(node, k);
    }
    
    int helper(MyTreeNode*& node, int k) {
        if (node == nullptr)
            return 0;
        if (node->left) {
            int cnt = node->left->count;
            if (k <= cnt) 
                return helper(node->left, k);
            else if (k > cnt + 1) 
                return helper(node->right, k - 1 - cnt);
            return node->val;
        } 
        else {
            if (k == 1) 
                return node->val;
            return helper(node->right, k - 1);
        }
    }
};