863. 二叉树中所有距离为 K 的结点

863. 二叉树中所有距离为 K 的结点

树上搜索

class Solution {
public:
    map<TreeNode *, TreeNode *> fa;
    vector<int> res;

    void getfa(TreeNode *root) {
        if (root->left) {
            fa[root->left] = root;
            getfa(root->left);
        }
        if (root->right) {
            fa[root->right] = root;
            getfa(root->right);
        }
    }

    void dfs(TreeNode *root, TreeNode *from, int k) {
        if (k == 0) {
            res.push_back(root->val);
            return;
        }

        if (root->left && root->left != from) {
            dfs(root->left, root, k - 1);
        }
        if (root->right && root->right != from) {
            dfs(root->right, root, k - 1);
        }
        if (fa.count(root) && fa[root] != from) {
            dfs(fa[root], root, k - 1);
        }
    }

    vector<int> distanceK(TreeNode *root, TreeNode *target, int k) {
        getfa(root);
        dfs(target,nullptr,k);
        return res;
    }
};