问题 23

题目描述

给定一个二叉树和一个节点，要求找出这个节点的下一个节点。节点的下一个节点是按照中序遍历次序的下一个节点。

函数签名

def inorder_successor(root: TreeNode, node: TreeNode) -> TreeNode:
    pass

输入

• root：二叉树的根节点
• node：要查找后继节点的节点

输出

• 返回节点的下一个节点，如果节点不存在，则返回 None。

示例

示例 1

Input:
        4
       / \
      2   6
     / \ / \
    1  3 5  7

node = 2

Output: 3

示例 2

Input:
        4
       / \
      2   6
     / \ / \
    1  3 5  7

node = 4

Output: 5

题解

本题需要在中序遍历的顺序中查找一个节点的后继节点。可以分为两种情况来讨论：

1. 如果当前节点的右子树存在，那么右子树中的最左边的节点就是当前节点的后继节点；
2. 如果当前节点的右子树不存在，那么需要往上找到一个节点，它还没有处理过它的右子树的祖先节点，那么这个祖先节点就是当前节点的后继。

根据上述思路，我们可以得到以下代码：

class TreeNode:
    def __init__(self, val=None, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

def inorder_successor(root: TreeNode, node: TreeNode) -> TreeNode:
    if node.right:
        node = node.right
        while node.left:
            node = node.left
        return node
    
    parent = None
    while root:
        if node.val < root.val:
            parent = root
            root = root.left
        elif node.val > root.val:
            root = root.right
        else:
            break
    return parent

首先，如果当前节点有右子树，那么就再右子树中找到最左边的节点。代码中通过一个 while 循环来遍历右子树中的节点，直到找到最左边的节点为止。

如果当前节点没有右子树，我们就需要往上找。我们可以定义一个 parent 变量来保存最后一次访问的大于当前节点值的节点，这个节点将会是当前节点的后继。如果当前节点的值小于当前节点的值，则说明它的后继节点一定在它父节点的右子树中，所以我们更新 parent 为当前节点。如果当前节点的值大于当前节点的值，则说明它的后继节点一定在它的左子树中，我们就递归地向下找到当前节点。当找到当前节点时，直接返回 parent 就是它的后继节点了。

复杂度分析

该算法时间复杂度为 $O(h)$，其中 $h$ 为二叉树的高度。空间复杂度为 $O(1)$，因为算法只使用了常数级的额外空间（递归栈空间不计入总空间）。