Description

Submission

/**
 * 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 {
    vector<int> rets;
    TreeNode* target;
    int k;

    void subDfs(TreeNode* node, int distance) {
        if(!node) return;
        if(distance < 0) return;
        if(distance == 0) {
            rets.push_back(node->val);
        }
        subDfs(node->left, distance - 1);
        subDfs(node->right, distance - 1);
    }

    int dfs(TreeNode* node) {
        if(!node) return 0;

        if(node == target) {
            subDfs(node, k);
            return 1;   // the height between the subroot and node 
        }
        int val;
        if(val = dfs(node->left)) {
            int d = k - val - 1;
            if(d == -1) {
                rets.push_back(node->val);
            } else {
                subDfs(node->right, d);
            }
            
            return val + 1;
        }

        if(val = dfs(node->right)) {
            int d = k - val - 1;
            if(d == -1) {
                rets.push_back(node->val);
            } else {
                subDfs(node->left, d);
            }
            return val + 1;
        }

        return 0;
    }


public:
    vector<int> distanceK(TreeNode* root, TreeNode* target, int k) {
        this->target = target;
        this->k = k;

        dfs(root);

        return rets;
    }
};