能迭代就不要递归——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;
}
};
需要注意的一点是,对于在栈中的结点,需要标记这个结点的左结点是否已经入栈,也就是说,当我们访问一个栈顶结点,如果它的左子结点还没有入栈,就改变标志位,然后左结点入栈;反之,就把栈顶结点弹出栈,再把该结点的右子结点入栈。因此栈内元素是结点和标志位组成的数据对。