JZ57 二叉树的下一个结点 python

JZ57 二叉树的下一个结点

book method
画图找规律，在中序遍历中寻找该节点的下一个节点分如下三类讨论：

1. 节点有右子树
   找到其右子树的最左子节点为中序遍历的下一节点
2. 节点无右子树，且为父节点的左节点
   其父节点为中序遍历的下一节点
3. 节点无右子树，且为父节点的右节点
   往上找父节点，直到该节点位于某个父节点的左子树，则这个父节点为中序遍历的下一节点

错误记录

1. 未考虑情况3中找不到所需父节点的情况（即为最后一个节点）{8,6,10,5,7,9,11},11
2. 未考虑根节点为中序遍历最后一个节点的情况{5,4,#,3,#,2},5

class TreeLinkNode:
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None
        self.next = None #指向父节点
class Solution:
    def GetNext(self, pNode):
        if pNode.right:
            # 节点有右子树
            p = pNode.right
            while p:
                if p.left:
                    p = p.left
                else:
                    return p
        else:
            # 节点无右子树
            if pNode.next:
                parent = pNode.next
                if pNode == parent.left:
                    # 节点为父节点的左节点
                    return parent
                else:
                    # 节点为父节点的右节点
                    while parent.next and parent != parent.next.left:
                        parent = parent.next
                return parent.next

优化

class Solution:
    def GetNext(self, pNode):
        # write code here
        # 此节点有右子树
        if pNode.right:
            temp = pNode.right
            while temp:
                result = temp
                temp = temp.left
            return result
        # 此节点没有右子树
        while pNode:
            if pNode.next:
                if pNode.next.left == pNode:
                    return pNode.next
                pNode = pNode.next
            else:
                return None

注意特殊二叉树的样例