[LeetCode]面试题04_06_后继者

题目链接：面试题04_06_后继者

题目描述

设计一个算法，找出二叉搜索树中指定节点的“下一个”节点（也即中序后继）。
如果指定节点没有对应的“下一个”节点，则返回null。
示例1:

输入: root = [2,1,3], p = 1
  2
 / \
1   3
输出: 2

示例2：

输入: root = [5,3,6,2,4,null,null,1], p = 6
      5
     / \
    3   6
   / \
  2   4
 /   
1

解题思路

• 中序遍历并记录：
  这道题考察的就是中序遍历，一种朴素的做法就是，中序遍历该二叉树，并且记录遍历的序列；然后遍历该序列，寻找给定节点的下一个节点。
  代码如下：

#include <iostream>
#include <vector>
#include <stack>
using namespace std;
struct TreeNode
{
    int val;
    TreeNode *left;
    TreeNode *right;
    TreeNode(int x) : val(x), left(NULL), right(NULL) {}
};

class Solution
{
public:
    TreeNode *inorderSuccessor(TreeNode *root, TreeNode *p)
    {
        vector<TreeNode *> seq; //记录中序遍历的序列
        stack<TreeNode *> st;   //非递归中序遍历所用到的栈

        TreeNode *r = root;
        //中序遍历
        while (!st.empty() || r != nullptr)
        {
            while (r != nullptr)
            {
                st.push(r);
                r = r->left; //将左子树一直入栈
            }
            r = st.top(); //弹栈
            st.pop();
            seq.push_back(r);
            r = r->right; //压入右子树
        }
        int i;
        for (i = 0; i < seq.size(); i++) //寻找节点p的下一个节点
        {
            if (seq[i]->val == p->val)
            {
                break;
            }
        }
        if (i == seq.size() - 1)
            return nullptr;
        else
            return seq[i + 1];
    }
};

时间复杂度为：$O(n)$
空间复杂度为：$O(n)$

• 记录前驱节点的中序遍历：事实上我们并不需要将整棵树都完整遍历下来，因为我们要求的只是节点p的后继节点。可以在遍历过程中记录前驱节点，如果在中序遍历过程中某节点的前驱节点是p，说明p的后继节点就是该节点，直接返回该节点即可。

代码如下：

#include <iostream>
#include <vector>
#include <stack>
using namespace std;
struct TreeNode
{
    int val;
    TreeNode *left;
    TreeNode *right;
    TreeNode(int x) : val(x), left(NULL), right(NULL) {}
};

class Solution
{
public:
    TreeNode *inorderSuccessor(TreeNode *root, TreeNode *p)
    {
        stack<TreeNode *> st;
        TreeNode *r = root;
        TreeNode *pre = nullptr;//记录前驱节点
        //中序遍历
        while (!st.empty() || r != nullptr)
        {
            while (r != nullptr)
            {
                st.push(r);
                r = r->left;
            }
            r = st.top();
            st.pop();
            if (pre == p)//如果某节点的前驱节点是p，则返回该节点
            {
                return r;
            }
            pre = r;//更新前驱节点为当前节点
            r = r->right;
        }
        //返回空
        return nullptr;
    }
};

空间复杂度：$O(1)$。
时间复杂度仍然是：$O(n)$，但平均时间复杂度要优于前一种算法。