[面试题 04.06][中等][DFS] 后继者

题目描述

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

解题思路

中序遍历的下一个元素，5大解法汇总！

由于本题中的树是二叉搜索树, 那么其中序遍历是单调递增的, 依照这个思路, 简化查找过程.

• 如果结点 p 的值大于等于 root 的值，说明 p 的后继结点在 root 右子树中，那么就递归到右子树中查找

• 如果结点 p 的值小于 root 的值，说明 p 在 root 左子树中，而它的后继结点有两种可能，要么也在左子树中，要么就是 root

  • 如果左子树中找到了后继结点，那就直接返回答案

  • 如果左子树中没有找到后继结点，那就说明 p 的右儿子为空，那么 root 就是它的后继结点

搜索的过程可以分为递归和非递归两种方式.

递归

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

class Solution:
    def inorderSuccessor(self, root: TreeNode, p: TreeNode) -> TreeNode:
        if root is None:
            return

        if p.val >= root.val:
            return self.inorderSuccessor(root.right, p)
        else:
            node = self.inorderSuccessor(root.left, p)
            return root if node is None else node

非递归

一图看懂二叉搜索树后继查找(非递归实现)

class Solution:
    def inorderSuccessor(self, root: TreeNode, p: TreeNode) -> TreeNode:
        # 首先向右查找, 如果有右子树, 结果肯定在这里面
        if p.right is not None:
            current = p.right
            while current is not None and current.left is not None:
                current = current.left
            return current

        # 右子树中没有, 从根节点出发, 寻找`p`结点的位置
        stack = []
        current = root
        while current is not p:
            stack.append(current)
            current = current.right if p.val > current.val else current.left

        while stack and stack[-1].left is not current:
            current = stack.pop(-1)

        if stack:
            return stack[-1]
        return None