劍指offer：二叉樹(shù)的下一個(gè)節(jié)點(diǎn)

題目描述
給定一個(gè)二叉樹(shù)和其中的一個(gè)結(jié)點(diǎn)，請(qǐng)找出中序遍歷順序的下一個(gè)結(jié)點(diǎn)并且返回。注意，樹(shù)中的結(jié)點(diǎn)不僅包含左右子結(jié)點(diǎn)，同時(shí)包含指向父結(jié)點(diǎn)的指針。

class TreeLinkNode:
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None
        self.parent = None

class Solution:
    """
    給定節(jié)點(diǎn)，求中序遍歷下該節(jié)點(diǎn)的下一個(gè)節(jié)點(diǎn)
    """
    def GetNext(self, pNode):
        if not pNode:
            return None

        # 如果給定節(jié)點(diǎn)存在右子樹(shù)，則下一個(gè)節(jié)點(diǎn)就是右子樹(shù)的最左葉子節(jié)點(diǎn)
        if pNode.right:
            temp = pNode.right
            while temp.left:
                temp = temp.left
            return temp

        if pNode.parent:
            # 如果給定節(jié)點(diǎn)不存在右子樹(shù)，則找到第一個(gè)在它右邊的祖先節(jié)點(diǎn)
            # 例如，給定節(jié)點(diǎn)是父節(jié)點(diǎn)的左子節(jié)點(diǎn)，則下一個(gè)節(jié)點(diǎn)就是父節(jié)點(diǎn)
            current = pNode
            parent = pNode.parent
            # 如果給定節(jié)點(diǎn)是父節(jié)點(diǎn)的右子節(jié)點(diǎn)，則需要繼續(xù)往上找，找到第一個(gè)祖先節(jié)點(diǎn)，
            # 使得該祖先節(jié)點(diǎn)的是其父節(jié)點(diǎn)的左子節(jié)點(diǎn)，則下一個(gè)節(jié)點(diǎn)就是該祖先節(jié)點(diǎn)的父節(jié)點(diǎn)(因?yàn)榻o定
            # 節(jié)點(diǎn)是該祖先節(jié)點(diǎn)的父節(jié)點(diǎn)左邊最近的子節(jié)點(diǎn))
            while parent and current == parent.right:
                current = parent
                parent = parent.parent
            return parent

        # 給定節(jié)點(diǎn)已是最右節(jié)點(diǎn)，不存在下一個(gè)節(jié)點(diǎn)
        return None