题目
给定一棵二叉树的根节点 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);
}
}
评论区