leetcode 863. All Nodes Distance K in Binary Tree

We are given a binary tree (with root node root), a target node, and an integer value K.

Return a list of the values of all nodes that have a distance K from the target node. The answer can be returned in any order.

Example 1:

Input: root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, K = 2
Output: [7,4,1]
Explanation: 
The nodes that are a distance 2 from the target node (with value 5)
have values 7, 4, and 1.

1

Note that the inputs "root" and "target" are actually TreeNodes.
The descriptions of the inputs above are just serializations of these objects.

Note:

The given tree is non-empty.
Each node in the tree has unique values 0 <= node.val <= 500.
The target node is a node in the tree.
0 <= K <= 1000.

思路:将二叉树转换为无向图,然后进行dfs记录深度。 注意,对无向图进行遍历的时候要记录父节点,又不会有重复。

/**
 * 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:
    map<int, vector<int>> edge;
    vector<int> ans;
    void dfs(TreeNode* root) {
        if (root == nullptr) return;
        if (root->left != nullptr) {
            dfs(root->left);
            edge[root->val].push_back(root->left->val);
            edge[root->left->val].push_back(root->val);
        }
        if (root->right != nullptr) {
            dfs(root->right); 
            edge[root->val].push_back(root->right->val);
            edge[root->right->val].push_back(root->val);
              
        }
    }
    void dfs2(int u, int fa, int dep, int K) {
        if (dep == K) {
            ans.push_back(u);
            return;
        }
        for (int i = 0; i < edge[u].size(); ++i) {
            int v = edge[u][i];
            if (v == fa) continue;
            dfs2(v, u, dep+1, K);
        }
    }
    vector<int> distanceK(TreeNode* root, TreeNode* target, int K) {
        dfs(root);
        dfs2(target->val, -1, 0, K);
        return ans;
    }
};
posted on 2018-07-01 11:44  Beserious  阅读(669)  评论(0编辑  收藏  举报