Given a binary tree (with root), a target, and an integer value K.
Returns a list of values for all nodes with a distance of K from the target node. The answers 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] Interpretation: The calculated node is a node with a distance of 2 from the target node (value 5). Values are 7, 4, and 1, respectively Note that the input "root" and "target" are actually nodes on the tree. The above input only describes the serialization of these objects.
Tips:
- The given tree is non empty and has at most K nodes.
- Each node on the tree has a unique value of 0 < = node. Val < = 500.
- The target node is the node on the tree.
- 0 <= K <= 1000.
C++
/** * Definition for a binary tree node. * struct TreeNode { * int val; * TreeNode *left; * TreeNode *right; * TreeNode(int x) : val(x), left(NULL), right(NULL) {} * }; */ class Solution { public: void dfs(TreeNode* root, map<TreeNode*,TreeNode*>& tmp) { if(root) { if(root->left) { tmp[root->left]=root; dfs(root->left,tmp); } if(root->right) { tmp[root->right]=root; dfs(root->right,tmp); } } } void dfs2(map<TreeNode*,TreeNode*>& tmp, map<TreeNode*,int>& flag, vector<int>& res, int K, TreeNode* node) { if(0==K) { if(node && flag[node]) { res.push_back(node->val); } return; } else { if(flag[node]) { flag[node]=0; if(tmp[node]) { dfs2(tmp,flag,res,K-1,tmp[node]); } if(node->left) { dfs2(tmp,flag,res,K-1,node->left); } if(node->right) { dfs2(tmp,flag,res,K-1,node->right); } } } } vector<int> distanceK(TreeNode* root, TreeNode* target, int K) { map<TreeNode*,TreeNode*> tmp; dfs(root,tmp); map<TreeNode*,int> flag; for(auto it:tmp) { flag[it.first]=1; } flag[root]=1; //Root node has no parent vector<int> res; dfs2(tmp,flag,res,K,target); return res; } };