863. All Nodes Distance K in Binary Tree

題目描述

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

解答

C++

/**
 * 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*, TreeNode*> parentMap;
    vector<int> distanceK(TreeNode* root, TreeNode* target, int k) {
        findParent(root, nullptr);
        vector<int> ans;
        
        queue<pair<TreeNode*, int>> frontier;
        frontier.push({target, 0});

        unordered_set<TreeNode*> visited;
        visited.insert(target);

        while (not frontier.empty()) {
            const auto [u, dist] = frontier.front();
            frontier.pop();
            cout << u->val << " ";
            if (dist == k) {
                ans.push_back(u->val);
            }
            if (u->left != nullptr and visited.count(u->left) == 0) {
                frontier.push({u->left, 1 + dist});
                visited.insert(u->left);
            }
            if (u->right != nullptr and visited.count(u->right) == 0) {
                frontier.push({u->right, 1 + dist});
                visited.insert(u->right);
            }
            if (parentMap[u] != nullptr and visited.count(parentMap[u]) == 0) {
                frontier.push({parentMap[u], 1 + dist});
                visited.insert(parentMap[u]);
            }
        }
        return ans;
    }
    void findParent(TreeNode* root, TreeNode* parent) {
        if (root == nullptr) {
            return ;
        }
        parentMap[root] = parent;
        findParent(root->left, root);
        findParent(root->right, root);
    }
};

Jerry Wu11 July, 2023

Python

class Solution:
    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:
        graph = defaultdict(list)

        def build(node, parent = None):
            if not node:
                return
            if parent:
                graph[node.val].append(parent.val)
                graph[parent.val].append(node.val)
            build(node.left, node)
            build(node.right, node)
        
        build(root)
        ans = deque([target.val])
        visited = set([target.val])
        for step in range(k):
            n = len(ans)
            for _ in range(n):
                node = ans.popleft()
                for child in graph[node]:
                    if child not in visited:
                        visited.add(child)
                        ans.append(child)
        return list(ans)

Yen-Chi ChenWed, Jul 12, 2023

Reference

回到題目列表