863. All Nodes Distance K in Binary Tree

Description

Given the root of a binary tree, the value of a target node target, and an integer k, return an array of the values of all nodes that have a distance k from the target node.

You can return the answer 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.

Example 2:

Input: root = [1], target = 1, k = 3
Output: []

Constraints:

• The number of nodes in the tree is in the range [1, 500].
• 0 <= Node.val <= 500
• All the values Node.val are unique.
• target is the value of one of the nodes in the tree.
• 0 <= k <= 1000

---

Solutions

Solution 1

Intuition

1. turn this tree to a undirected graph
2. dfs from target node, find k-distance node

Code

/**
 * 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<TreeNode*, vector<TreeNode*>> g;
    vector<int> res;

    void dfs(TreeNode* node) {
        if (node -> left) {
            g[node].push_back(node -> left);
            g[node -> left].push_back(node);
            dfs(node -> left);
        }
        if (node -> right) {
            g[node].push_back(node -> right);
            g[node -> right].push_back(node);
            dfs(node -> right);
        }
    }

    void find(TreeNode* node, TreeNode* father, int k) {
        if (!k) {
            res.push_back(node -> val);
        } else {
            for (auto son : g[node]) {
                if (son != father)
                    find(son, node, k - 1);
            }
        }
    }

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

Complexity

• Time: O(N)
• Space: O(N)

---