剑指offer - 二叉树的下一个结点 - python

题目描述：

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

思路：

二叉树的中序遍历顺序为左 -> 根 -> 右，那么需要首先判断给定的结点属于哪一种：

• 如果它存在右子树，说明它是子树的父结点，下一个节点在其右子树中。如果右子树存在左孩子结点，那么下一个就是左孩子结点，否则说明它就是叶子结点，直接返回
• 如果它不存在右子树，说明当前结点为叶子节点：
  • 如果当前结点为右叶子，那么下一个就是它的父结点，即它next域指向的结点
  • 如果它的父节点有左叶子，说明下一个节点就是父节点

---

AC代码

# -*- coding:utf-8 -*-
# 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 == None: 
        	return None
        	
        # 如果右子树存在，说明当前节点为子树根结点，下一个结点在右子树
        if pNode.right:
            # 更新结点
            pNode = pNode.right
            # 中序遍历：左-根-右
            # 如果左孩子结点存在则返回
            while pNode.left:
                pNode = pNode.left
            # 否则说明当前结点即叶节点，直接返回
            return pNode
        # 如果右子树不存在，那么当前结点为叶节点
        else:
            # 如果next域有值，不管它是左孩子还是右孩子，先找它的父节点
            while pNode.next:
                # 如果父节点的左孩子存在，那么说明它是右孩子结点，下一个节点即父节点
                if pNode == pNode.next.left:
                    return pNode.next
                # 否则说明它是左孩子结点，下一个结点即父节点
                pNode = pNode.next
        return None