二叉树的下一个结点

题目描述

给定一个二叉树和其中的一个结点，请找出中序遍历顺序的下一个结点并且返回。注意，树中的结点不仅包含左右子结点，同时包含指向父结点的指针。

代码

/*
struct TreeLinkNode {
    int val;
    struct TreeLinkNode *left;
    struct TreeLinkNode *right;
    struct TreeLinkNode *next;
    TreeLinkNode(int x) :val(x), left(NULL), right(NULL), next(NULL) {
        
    }
};
*/
class Solution {
public:
    TreeLinkNode* GetNext(TreeLinkNode* pNode)
    {
        if (pNode == NULL) {
            return NULL;
        }
        TreeLinkNode* cur = pNode;
        while (cur->next != NULL) {//先找根节点
            cur = cur->next;
        }
        int tag = 0;
        return dfs(cur, pNode, tag);
    }
    TreeLinkNode* dfs(TreeLinkNode* r, TreeLinkNode* p, int &tag) {
        if (r == NULL) {
            return NULL;
        }
        TreeLinkNode *left = dfs(r->left, p, tag);
        if (left != NULL) {//左子树找到了
            return left;
        }
        if (tag == 1) {//第一次tag为1时就是答案
            tag = 0;
            return r;
        }
        if (p == r || tag) {
            tag = 1;
        }
        return dfs(r->right, p, tag);
    }
};