解法一:中序遍历
在迭代法中序遍历的同时,用一个布尔值来记录是否找到了目标节点,找到了就将它设置为
true
,返回下一个节点即可。时间复杂度为$O(N)$,因为要用一个栈来保存顺着指向左子节点的指针的路径上的所有节点,因此空间复杂度为$O(h)$,h
为二叉树的高度。
class Solution {
public:
TreeNode* inorderSuccessor(TreeNode* root, TreeNode* p) {
stack<TreeNode*> stk;
TreeNode* cur = root;
bool flag = false;
while (cur || !stk.empty())
{
while (cur != nullptr)
{
stk.push(cur);
cur = cur->left;
}
cur = stk.top();
stk.pop();
if (flag) //若flag已经为true,则当前的cur就是目标节点的下一节点
{
break;
}
else if (cur == p) //找到目标节点就将flag置为true
{
flag = true;
}
cur = cur->right;
}
return cur;
}
};
解法二
从根节点开始遍历,将每个节点的值与目标节点值比较,因为在中序遍历中目标节点的下一节点值肯定比它要大,遍历时分两种情况讨论:
- 若目标节点值小于当前节点值,则走向它的左子树,并将结果指针指向当前节点;
- 若目标节点值大于当前节点值,则走向它的右子树。
该方法时间复杂度为$O(N)$,空间复杂度为$O(1)$。
class Solution {
public:
TreeNode* inorderSuccessor(TreeNode* root, TreeNode* p) {
TreeNode* cur = root;
TreeNode* res = nullptr;
while (cur != nullptr)
{
if (p->val < cur->val) //第一种情况
{
res = cur;
cur = cur->left;
}
else //第二种情况
{
cur = cur->right;
}
}
return res;
}
};