We are given a binary tree (with root node root), a target node, and an integer value K.
Return a list of the values of all nodes that have a distance K from the target node. The answer can be returned in any order.
Example 1:
Input: root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, K = 2
Output: [7,4,1]
Explanation:
The nodes that are a distance 2 from the target node (with value 5)
have values 7, 4, and 1.
这道题不能一个recursion写完。
要分两个recursion来写。
先把这个点找出来. 找的时候同时记录从parent to 这个node的路径上的所有点, 可以存在一个stack里面。
这个可以用一个recursion来做。
然后从stack里面一个一个取,从每个点开始找相应的点。
这里面你会发现最后一个点和其他的点不一样。
最后一个点你可以既找它的左孩子,又可以找它的右孩子。但其他点只能找左孩子右孩子中间的其中一个。
所以我这里面用了一个prev来记录上一个node,如果上一个node是当前是node的左孩子,我就只从当前node右子树去找。
class Solution {
public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
List<Integer> ans = new ArrayList<>();
Deque<TreeNode> deque = new ArrayDeque<>();
if (root == null || target == null) return ans;
findPath(root, deque, target);
TreeNode prev = null;
int dist = 0;
while (!deque.isEmpty()) {
TreeNode node = deque.pop();
if (dist == K) {
ans.add(node.val);
return ans;
}
if (node.left != prev) findAllNode(node.left, K - dist - 1, ans);
if (node.right != prev) findAllNode(node.right, K - dist - 1, ans);
prev = node;
dist++;
}
return ans;
}
public void findAllNode(TreeNode node, int level, List<Integer> ans) {
if (node == null || level < 0) return;
if (level == 0) {
ans.add(node.val);
return;
}
findAllNode(node.left, level - 1, ans);
findAllNode(node.right, level - 1, ans);
}
private boolean findPath(TreeNode node, Deque<TreeNode> deque, TreeNode target) {
if (node == null) return false;
deque.push(node);
if (node == target) return true;
if (findPath(node.left, deque, target)) return true;
if (findPath(node.right, deque, target)) return true;
deque.pop();
return false;
}
}
网友评论