能迭代就不要递归——Leetcode 230 二叉搜索树中第k小的元素

能迭代就不要递归——Leetcode 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 函数？

二叉搜索树是排好序的数据结构，因此我们只需要通过中序遍历，把所有元素按从小到大的顺序取出来，就可以找到第k小的元素了。而既然是要求第k小的元素，那么只需要遍历前k个元素就可以了。相应的代码如下：

/**
 * 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) 
    {
        inorderTraversal(root, k);    
        return ans;
    }
    void inorderTraversal(TreeNode* root, int &k)
    {
        if (!root) { return; }
        inorderTraversal(root->left, k);
        if (k == 0) { return; }
        --k;
        if (k == 0) { ans = root->val; }
        inorderTraversal(root->right, k);
    }
private:
    int ans;
};

如果是最原始的中序遍历，那么递归就可以实现了，但是如果我们想要在合适的时候结束递归，事情就变得没那么简单。因为我们需要从递归的最内层退出到最外层递归的入口。为了实现这一点，我们不得不引入一些判断条件，使得递归函数在满足这些条件的时候返回。因此，当k等于0的时候，说明我们已经找到第k小的结点了，需要返回，这就是引入的判断条件。

还有没有别的解决办法呢？很多递归问题都能用迭代实现（能迭代实现的问题一定能递归实现，但是能递归实现的问题不一定能迭代实现），这个问题可不可以用迭代实现？学过数据结构的人都知道，中序遍历是可以用迭代实现的。因此很容易就写出迭代实现的版本：

/**
 * 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<pair<TreeNode*, bool> > sta;
        sta.push(make_pair(root, true));
        while (!sta.empty())
        {
            if (sta.top().second)
            {
                sta.top().second = false;
                if (sta.top().first->left)
                {
                    sta.push(make_pair(sta.top().first->left, true));
                }           
            }
            else
            {
                TreeNode* top = sta.top().first;
                sta.pop();
                --k;
                if (k == 0) { return top->val; }
                if (top->right)
                {
                    sta.push(make_pair(top->right, true));
                }
            }
        }
        return NULL;
    }
};

需要注意的一点是，对于在栈中的结点，需要标记这个结点的左结点是否已经入栈，也就是说，当我们访问一个栈顶结点，如果它的左子结点还没有入栈，就改变标志位，然后左结点入栈；反之，就把栈顶结点弹出栈，再把该结点的右子结点入栈。因此栈内元素是结点和标志位组成的数据对。