[LeetCode] All Nodes Distance K in Binary Tree

想法：

先透過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;
    }
};