863

863. All Nodes Distance K in Binary Tree (Medium)

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.



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

 

Note:

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

Companies:
Amazon, Google

Related Topics:
Tree, Depth-first Search, Breadth-first Search

Solution 1. DFS + Level Order

Consider a simpler question: given root, find the nodes that have a distance K from the root node. We can simply use level order traversal to find the nodes at Kth level.

For this question, we need to consider the target's parents and the nodes in parents' substrees. We can use DFS to locate the target node, and if target node is found, return the distance to target node; otherwise return INT_MAX which means the current node is not target's parent and should be ignored.

For the target's parents, we use the level order traversal in the other subtree that doesn't contain target.

// OJ: https://leetcode.com/problems/all-nodes-distance-k-in-binary-tree/
// Author: github.com/lzl124631x
// Time: O(N)
// Space: O(N)
class Solution {
private:
    vector<int> ans;
    TreeNode *target;
    int K;
    int inorder(TreeNode *root) {
        if (!root) return INT_MAX;
        if (root == target) return 1;
        int left = inorder(root->left);
        if (left != INT_MAX) {
            if (left < K) levelOrder(root->right, K - left - 1);
            else if (left == K) ans.push_back(root->val);
            return left + 1;
        }
        int right = inorder(root->right);
        if (right != INT_MAX) {
            if (right < K) {
                levelOrder(root->left, K - right - 1);
            } else if (right == K) ans.push_back(root->val);
            return right + 1;
        }
        return INT_MAX;
    }
    void levelOrder(TreeNode *root, int K) {
        if (!root) return;
        queue<TreeNode*> q;
        q.push(root);
        int level = 0;
        while (q.size()) {
            int cnt = q.size();
            while (cnt--) {
                TreeNode *node = q.front();
                q.pop();
                if (level == K) ans.push_back(node->val);
                else {
                    if (node->left) q.push(node->left);
                    if (node->right) q.push(node->right);
                }
            }
            ++level;
        }
    }
public:
    vector<int> distanceK(TreeNode* root, TreeNode* target, int K) {
        this->target = target;
        this->K = K;
        levelOrder(target, K);
        inorder(root);
        return ans;
    }
};

Solution 2. DFS + BFS

DFS to build the graph of the tree, then BFS to find the nodes at level k.

// OJ: https://leetcode.com/problems/all-nodes-distance-k-in-binary-tree
// Author: github.com/lzl124631x
// Time: O(N)
// Space: O(N)
class Solution {
public:
    vector<int> distanceK(TreeNode* root, TreeNode* target, int k) {
        unordered_map<int, vector<int>> G;
        function<void(TreeNode*, TreeNode*)> dfs = [&](TreeNode *node, TreeNode *parent) {
            if (!node) return;
            if (parent) {
                G[parent->val].push_back(node->val);
                G[node->val].push_back(parent->val);
            }
            dfs(node->left, node);
            dfs(node->right, node);
        };
        dfs(root, nullptr);
        queue<int> q{{target->val}};
        while (q.size() && k--) {
            int cnt = q.size();
            while (cnt--) {
                int u = q.front();
                q.pop();
                for (int v : G[u]) {
                    if (G.count(v)) q.push(v);
                }
                G.erase(u);
            }
        }
        vector<int> ans;
        while (q.size()) {
            ans.push_back(q.front());
            q.pop();
        }
        return ans;
    }
};