想法:
先透過DFS算出disance between target and ancestors
在做一次DFS or BFS把距離的information傳遞下去給其他分支
偷懶直接用array存距離而不用hash table,畢竟只有題目說數字低於500
class Solution { public: int dis[501]; vector<int> ans; vector<int> distanceK(TreeNode* root, TreeNode* target, int K) { foundPath(root,target); DFS(root, -1, target, K); return ans; } void DFS(TreeNode* root, int pd, TreeNode* target, int K) { if(dis[root->val] == 0 && root != target) dis[root->val] = pd; if(dis[root->val] == K) ans.push_back(root->val); if(root->left) { if(dis[root->left->val] == 0) DFS(root->left,dis[root->val]+1, target, K); else DFS(root->left,dis[root->left->val], target, K); } if(root->right) { if(dis[root->right->val] == 0) DFS(root->right,dis[root->val]+1, target, K); else DFS(root->right,dis[root->right->val], target, K); } } //look for path between target and root, calculate disance each node on the path to target int foundPath(TreeNode* root, TreeNode* target) { if(!root) return -1; if(root == target) return 0; int rd = 0, ld = 0; if(root->left) { ld = foundPath(root->left, target); if(ld >= 0) { dis[root->val] = ld+1; return ld+1; } } if(root->right) { rd = foundPath(root->right, target); if(rd >= 0) { dis[root->val] = rd+1; return rd+1; } } return -1; } };