原題

Given a binary search tree and a node in it, find the in-order successor of that node in the BST.

Given tree = [2,1] and node = 1:

   2
  /
 1

return node 2.
Given tree = [2,1,3] and node = 2:

    2
   / \
  1   3

return node 3.

解題思路

• successor變量記錄root的父親結點，當while循環(huán)結束時，如果原本的root = None或者p不存在與BST中(那么此刻root = None)，都返回None
• 找到p之后，如果p沒有右兒子，則第一個比它大的數(shù)字就是剛剛記錄的successor
• 找到p之后，如果有右兒子，則找到右子樹中的最左邊的值(最小值)

完整代碼

# Definition for a binary tree node.
# class TreeNode(object):
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution(object):
    """
    @param root <TreeNode>: The root of the BST.
    @param p <TreeNode>: You need find the successor node of p.
    @return <TreeNode>: Successor of p.
    """
    def inorderSuccessor(self, root, p):
        successor = None
        while root != None and root.val != p.val:
            if root.val > p.val:
                successor = root
                root = root.left
            else:
                root = root.right
        
        # 原本的root = None 或者 p不存在與BST中，此刻root = None
        if root == None:
            return None
                
        # 找到p之后，如果p沒有右兒子，則第一個比它大的數(shù)字就是剛剛記錄的successor
        if root.right == None:
            return successor
        
        # 找到p之后，如果有右兒子，則找到右子樹中的最左邊的值（最小值）
        root = root.right
        while root.left != None:
            root = root.left
        
        return root