863. All Nodes Distance K in Binary Tree

(1) 关键找到target node的 所有root, 包括自身；

(2) 排除target所在子树， 找到距离为 (k-i) > 0 的node list。其中i为root到target的距离。

 

class Solution {
    
   
    //find path
    vector<TreeNode*> findTarget(TreeNode* root, TreeNode* target, vector<TreeNode*> path) {
        if(root == target)
            return path;
        if(root == NULL) {
            vector<TreeNode*> ret;
            return ret;
        }
        
        path.push_back(root);
        
        vector<TreeNode*> pathL = findTarget(root->left,target,path);
        vector<TreeNode*> pathR = findTarget(root->right,target,path);
        
        if(pathL.size() > 0) 
            return pathL;
        else
            return pathR;
        
    }
    

   //add node
    void addNode(TreeNode* root, TreeNode* non, int k, vector<int> &ret) {
        
        if(k<0)
            return;
        
        stack<TreeNode*> stk;
        stk.push(root);
    
        int count = 0;
        while(!stk.empty()) {
            
            if(count == k) {
                while(!stk.empty()) {
                    ret.push_back(stk.top()->val);
                    stk.pop();
                }
                return;
            }
            
            vector<TreeNode*> v;
            while(!stk.empty()) {
                TreeNode* tmp = stk.top();
                stk.pop();
                if(tmp->left && tmp->left != non)
                    v.push_back(tmp->left);
                if(tmp->right && tmp->right != non)
                    v.push_back(tmp->right);
            }
            
            for(int i=0;i<v.size();i++) {
                stk.push(v[i]);
            }
            
            count++;
        }
        return;
    }
public:
    vector<int> distanceK(TreeNode* root, TreeNode* target, int K) {
        vector<int> ret;
        
        if(root == NULL || target == NULL || K<0) return ret;
        
        vector<TreeNode*> path;
        path = findTarget(root,target,path);    
        path.push_back(target);
        int k=K;
        for(int i=path.size()-1;i>=0;i--) {
           
            TreeNode* tmp = NULL;
            if(i+1 < path.size())
                tmp = path[i+1];
            addNode(path[i],tmp,k,ret);
            k--;
        }
        
        return ret;
    }
};