863. All Nodes Distance K in Binary Tree

时间：2021-03-09 13:09:46 收藏：0 阅读：0

问题：

给定一棵二叉树。

求给定节点target开始，距离K的所有节点。

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.

Note that the inputs "root" and "target" are actually TreeNodes.
The descriptions of the inputs above are just serializations of these objects.
 

Note:
The given tree is non-empty.
Each node in the tree has unique values 0 <= node.val <= 500.
The target node is a node in the tree.
0 <= K <= 1000.

技术图片

解法：DFS+BFS

思想：

首先找到target节点，
- 由于需要其父节点路径，因此使用->DFS
对当前target节点：
- 向下：遍历找到距离层次为K的所有子节点->BFS
- 向上：每上一层父节点，K--，从其父节点再向下：遍历找到距离层次为K的所有子节点->BFS
  - 如果K==0，那么只需返回父节点即可，无需继续从该父节点继续向下展开。

?? 注意：对于每个父节点，需要知道的信息为自己应该向下遍历的K应该是多少。

因此DFS返回K值。

DFS的状态为：
- 当前节点root
BFS的状态为：
- 当前节点root，展开层数K

代码参考：

 1 /**
 2  * Definition for a binary tree node.
 3  * struct TreeNode {
 4  *     int val;
 5  *     TreeNode *left;
 6  *     TreeNode *right;
 7  *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 8  * };
 9  */
10 class Solution {
11 public:
12     void bfs(vector<int>& res, TreeNode* root, int K) {
13         queue<TreeNode*> q;
14         q.push(root);
15         if(K<0) return;
16         while(!q.empty()) {
17             int sz = q.size();
18             for(int i=0; i<sz; i++) {
19                 TreeNode* cur = q.front();
20                 q.pop();
21                 if(K==0) res.push_back(cur->val);
22                 if(cur->left) q.push(cur->left);
23                 if(cur->right) q.push(cur->right);
24             }
25             if(K==0) return;
26             K--;
27         }
28         return;
29     }
30     int dfs(vector<int>& res, TreeNode* root, TreeNode* target, int K) {
31         if(!root) return INT_MAX;
32         int left = INT_MAX, right = INT_MAX;
33         if(root->val == target->val) {
34             bfs(res, root, K);
35             return K-1;
36         } else {
37             if(root->left) left = dfs(res, root->left, target, K);
38             if(left==INT_MAX && root->right) right = dfs(res, root->right, target, K);
39         }
40         if(left!=INT_MAX) {
41             if(left==0)res.push_back(root->val);
42             else if(root->right) bfs(res, root->right, left-1);
43             return left-1;
44         } else if(right!=INT_MAX) {
45             if(right==0)res.push_back(root->val);
46             else if(root->left)bfs(res, root->left, right-1);
47             return right-1;
48         }
49         return INT_MAX;
50     }
51     vector<int> distanceK(TreeNode* root, TreeNode* target, int K) {
52         vector<int> res;
53         dfs(res, root, target, K);
54         return res;
55     }
56 };