【树-中等】面试题 04.06. 后继者

【题目】
设计一个算法，找出二叉搜索树中指定节点的“下一个”节点（也即中序后继）。

如果指定节点没有对应的“下一个”节点，则返回null。

示例 1:

输入: root = [2,1,3], p = 1

2
/
1 3

输出: 2

示例 2:

输入: root = [5,3,6,2,4,null,null,1], p = 6

  5
 / \
3   6

/
2 4
/
1

输出: null
【代码】

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

class Solution:
    def dfs(self,root,p):
        if not root:
            return 
        self.dfs(root.left,p)
        if self.cnt==1:
            self.ans=root
            self.cnt+=1
        if root.val==p.val:
            self.cnt+=1
        self.dfs(root.right,p)
    def inorderSuccessor(self, root: TreeNode, p: TreeNode) -> TreeNode:
        self.ans=None
        self.cnt=0
        self.dfs(root,p)
        return self.ans

【方法2】

class Solution:
    def inorderSuccessor(self, root: TreeNode, p: TreeNode) -> TreeNode:
        pre=None
        while root.val!=p.val:
            if root.val<p.val:
                root=root.right
            else:
                pre=root
                root=root.left
        if root.right:
            root=root.right
            while root.left:
                root=root.left
            return root
        else:
            return pre