【迭代, BFS】找到二叉树中最近的右侧节点

题目

1602. 找到二叉树中最近的右侧节点

---

给定一棵二叉树的根节点 root 和树中的一个节点 u ，返回与 u 所在层中距离最近的右侧节点，当 u 是所在层中最右侧的节点，返回 null 。

示例 1：

输入：root = [1,2,3,null,4,5,6], u = 4
输出：5
解释：节点 4 所在层中，最近的右侧节点是节点 5。

示例 2：

输入：root = [3,null,4,2], u = 2
输出：null
解释：2 的右侧没有节点。

示例 3：

输入：root = [1], u = 1
输出：null

示例 4：

输入：root = [3,4,2,null,null,null,1], u = 4
输出：2

提示:

• 树中节点个数的范围是 [1, 10^5] 。
• 1 <= Node.val <= 10^5
• 树中所有节点的值是唯一的。
• u 是以 root 为根的二叉树的一个节点。

解题

方法一： 迭代 BFS 层序遍历

思路

使用队列辅助迭代进行层序遍历，如果遍历到的节点与 u 相同：

• 如果当前节点是这一层中最后一个节点，那 u 就是这层中最右的节点，返回 null
• 否则队列的队头就是 u 的右侧节点，返回该节点。

代码

class Solution {
    public TreeNode findNearestRightNode(TreeNode root, TreeNode u) {
        Queue<TreeNode> nodeQueue = new LinkedList<>();
        nodeQueue.offer(root);
        while (!nodeQueue.isEmpty()) {
            int size = nodeQueue.size();
            for (int i = 0; i < size; i++) {
                root = nodeQueue.poll();
                if (root.val == u.val) {
                    return i == size - 1 ? null : nodeQueue.peek();
                }
                if (root.left != null) nodeQueue.offer(root.left);
                if (root.right != null) nodeQueue.offer(root.right);
            }
        }
        return null;
    }
}

方法二： 递归 层序遍历

思路

使用递归的方法进行层序遍历，只不过把存储遍历结果的数据结构换成栈，如果遍历到某个节点，发现这一层遍历结果的栈顶元素是 u，那么当前遍历到的节点就是要求的 u 所在层的最近右侧节点，记录下这个节点，在主方法中返回。

代码

class Solution {
    private List<Deque<Integer>> helper = new ArrayList<>();
    private int uVal;
    private TreeNode ans;

    public TreeNode findNearestRightNode(TreeNode root, TreeNode u) {
        uVal = u.val;
        inorderTraversal(root, 0);
        return ans;
    }

    public void inorderTraversal(TreeNode root, int level) {
        if (root == null) return;
        if (level == helper.size()) {
            helper.add(new LinkedList<>());
        }

        Deque<Integer> currLevel = helper.get(level);
        if (!currLevel.isEmpty() && currLevel.peek() == uVal) ans = root; 
        currLevel.push(root.val);
        inorderTraversal(root.left, level + 1);
        inorderTraversal(root.right, level + 1);
    }
}

优化

剪枝：如果找到了 u 就记录下这一层的层数，把这一层以后的递归全部停止。

class Solution {
    private List<Deque<Integer>> helper = new ArrayList<>();
    private int uVal, stop = Integer.MAX_VALUE;
    private TreeNode ans;

    public TreeNode findNearestRightNode(TreeNode root, TreeNode u) {
        uVal = u.val;
        inorderTraversal(root, 0);
        return ans;
    }

    public void inorderTraversal(TreeNode root, int level) {
        if (root == null) return;
        if (level == helper.size()) {
            helper.add(new LinkedList<>());
        }

        Deque<Integer> currLevel = helper.get(level);
        if (!currLevel.isEmpty() && currLevel.peek() == uVal) {
            ans = root;
            stop = level;
        }
        currLevel.push(root.val);
        if (level >= stop) return;
        inorderTraversal(root.left, level + 1);
        inorderTraversal(root.right, level + 1);
    }
}