[LeetCode] 863 二叉树中所有距离为K的节点

863. 二叉树中所有距离为 K 的结点 - 力扣（LeetCode） (leetcode-cn.com)

思想就是用dfs，不过要做一些额外的处理。

1. 因为从target开始遍历，所以要单独存储节点之间的父子关系，可以用unordered_map。这道题里数字都不相同，所以当前节点的值作为key即可。从头遍历一遍树，存到哈希表中。

2.从target开始的dfs，向三个方向深搜：左右孩子和父节点。可以设一个depth参数记录遍历深度，好作为终止条件。需要引入一个TreeNode* from参数记录上一次遍历的节点，来避免深搜走回头路。

/**
 * 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 {
public:
    unordered_map<int, TreeNode*> parents; //当前节点值 父节点
    vector<int> ans;

    //为每个节点找到父节点
    void setParent(TreeNode* node) {
        if(node->left) {
            parents[node->left->val] = node;
            setParent(node->left);
        }
        if(node->right) {
            parents[node->right->val] = node;
            setParent(node->right);
        }
    }

    void dfs(TreeNode* node, TreeNode* from, int k, int depth) {
        if(node == nullptr) return;
        //终止条件
        if(depth == k) {
            ans.push_back(node->val);
            return;
        }
        
        if(node->left != from) {
            dfs(node->left, node, k, depth + 1);
        }
        if(node->right != from) {
            dfs(node->right, node, k, depth + 1);
        }
        if(parents[node->val] != from) {
            dfs(parents[node->val], node, k, depth + 1);
        }
    }


    vector<int> distanceK(TreeNode* root, TreeNode* target, int k) {
        setParent(root);
        dfs(target, nullptr, k, 0);
        return ans;
    }
};