Nodes at Distance K

• Time : $O(n)$
• Space : $O(h)$

/**
 * Definition for binary tree
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
void dfs2(TreeNode* node, int d, vector<int>& res, TreeNode* ignore) {
    if(!node) return;
    if(d == 0) res.push_back(node->val);
    else {
        if(node->left and node->left != ignore) dfs2(node->left, d - 1, res, ignore);
        if(node->right and node->right != ignore) dfs2(node->right, d - 1, res, ignore);
    }
}
void dfs(TreeNode* node, int t, int d, vector<int>& res, vector<TreeNode*>& st) {
    if(!node) return;
    if(node->val == t) {
        TreeNode* ignore = node;
        dfs2(node->left, d - 1, res, node);
        dfs2(node->right, d - 1, res, node);
        for(int i = st.size() - 1, dis = d - 1; i >= 0 and dis >= 0; i--, dis--) {
            dfs2(st[i], dis, res, ignore);
            ignore = st[i];
        }
        return;
    }
    st.push_back(node);
    dfs(node->left, t, d, res, st);
    dfs(node->right, t, d, res, st);
    st.pop_back();
}
vector<int> Solution::distanceK(TreeNode* A, int B, int C) {
    vector<int> res;
    vector<TreeNode*> st;
    dfs(A,B,C,res,st);
    return res;
}