题解 | #距离是K的二叉树节点#

#include <unordered_map>
#include <vector>
class Solution {
private:
    int k_{};
    int target_{};
    vector<int> res_;
    unordered_map<int, TreeNode*> parents_;
public:
    void findParents(TreeNode* root) {
        if (root->left) {
            parents_[root->left->val] = root;
            findParents(root->left);
        }
        if (root->right) {
            parents_[root->right->val] = root;
            findParents(root->right);
        }
    }
    TreeNode* findNode(TreeNode* root) {
        if (!root) return nullptr;
        if (root->val == target_) return root;
        auto l = findNode(root->left);
        if (l) return l;
        auto r = findNode(root->right);
        return r;
    }
    void func(TreeNode* root, TreeNode* parent, int depth) {
        if (!root) return;
        if (depth == k_) {
            res_.push_back(root->val);
            return;
        }
        if (root->left != parent) func(root->left, root, depth + 1);
        if (root->right != parent) func(root->right, root, depth + 1);
        if (parents_[root->val] != parent) func(parents_[root->val], root, depth + 1);
    }
    vector<int> distanceKnodes(TreeNode* root, int target, int k) {
        k_ = k;
        target_ = target;
        findParents(root);
        auto tarNode = findNode(root);
        func(tarNode, nullptr, 0);
        return res_;
    }
};