面试题 04.06. 后继者

题目

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

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

示例 1:

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

2
/
1 3

输出: 2

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

输出: null

来源：力扣（LeetCode）
链接：https://leetcode.cn/problems/successor-lcci
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

思路

思路写在注释里了。

代码

递归 O(n)空间

class Solution {
    public TreeNode inorderSuccessor(TreeNode root, TreeNode p) {
        if (root == null) {
            return null;
        }
        // 如果当前值不大于p的值，说明要找的节点一定在右子树
        if (root.val <= p.val) {
            return inorderSuccessor(root.right, p);
        }
        // 接下来判断是否可能是根节点
        // 先在左子树中找
        TreeNode left = inorderSuccessor(root.left, p);
        // 如果没找到返回根节点
        if (left == null) {
            return root;
        }
        // 找到了，返回在左子树中找到的节点
        return left;
    }
}

迭代 O(n)空间

常规迭代的中序遍历，没有利用二叉搜索树的性质。记录好上一个节点和p作比较即可。不推荐

class Solution {
    public TreeNode inorderSuccessor(TreeNode root, TreeNode p) {
        Deque<TreeNode> stack = new ArrayDeque<>();
        TreeNode pp = root;
        TreeNode pre = null;
        while (pp != null || !stack.isEmpty()) {
            if (pp != null) {
                stack.push(pp);
                pp = pp.left;
            } else {
                TreeNode cur = stack.poll();
                if (pre == p) {
                    return cur;
                }
                pre = cur;
                pp = cur.right;
            }
        }
        return null;
    }
}

迭代 O(1)空间

学习一下：利用二叉搜索树的性质后，O(n)空间可以省略。可以看做是树上的二分。

class Solution {
    public TreeNode inorderSuccessor(TreeNode root, TreeNode p) {
        TreeNode cur = root;
        TreeNode ans = null;
        while (cur != null) {
            // 同递归写法中逻辑，如果当前值不大于p的值，说明要找的节点一定在右子树
            if (cur.val <= p.val) {
                cur = cur.right;
            } else {
                // 不在右子树，那么可能是根节点，先暂时将ans置为根节点，然后到左子树看
                ans = cur;
                cur = cur.left;
            }
        }
        return ans;
    }
}