面试题 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

分析： 二叉搜索树的中序遍历，可以用递归或者非递归方式

方法一： 递归法

class Solution:
    def inorderSuccessor(self, root: TreeNode, p: TreeNode) -> TreeNode:
        def LDR(root, p, find):
            if not root:
                return find, None
            find, find_val = LDR(root.left, p, find)
            if find: 
                if find_val == None: 
                    find_val = root
                return find, find_val
            if p.val == root.val:
                find = True
            find, find_val = LDR(root.right, p, find)
            return find, find_val
        find, find_val = LDR(root, p, False)
        return find_val

方法二： 非递归法

class Solution:
    def inorderSuccessor(self, root: TreeNode, p: TreeNode) -> TreeNode:
        queue = []
        find = False
        while queue or root:
            while root:
                queue.append(root)
                root = root.left
            if find:
                if queue:
                    return queue.pop()
                else:
                    return None
            if queue:
                root = queue.pop()
                if p.val == root.val:
                    find = True
            
            root = root.right
        if queue:
            return queue.pop()
        else:
            return None