距目标 K 距离内的节点总和

给定一棵二叉树，一个目标节点和一个正整数K ，任务是找到距离目标节点K内所有节点的总和（包括总和中目标节点的值）。

例子：

Input: target = 9, K = 1,  
Binary Tree =             1
                                /  \
                             2     9
                           /      /   \
                        4      5     7
                      /  \           /  \
                   8    19      20   11
                 /      /   \
             30     40   50
Output: 22
Explanation: Nodes within distance 1 from 9 is 9 + 5 + 7 + 1 = 22

Input: target = 40,  K = 2,  
Binary Tree =             1
                                /  \
                             2     9
                           /      /   \
                        4      5     7
                      /  \           /  \
                   8    19      20   11
                 /      /   \
             30     40   50
Output: 113
Explanation: Nodes within distance 2 from 40 is
40 + 19 + 50 + 4 = 113

方法：这个问题可以基于以下思想使用散列和深度优先搜索来解决：

Use a data structure to store the parent of each node. Now utilise that data structure to perform a DFS traversal from target and calculate the sum of all the nodes within K distance from that node.

按照下面提到的步骤来实施该方法：

• 创建一个哈希表（比如par ）来存储每个节点的父节点。
• 执行 DFS 并存储每个节点的父节点。
• 现在在树中找到目标。
• 创建一个哈希表来标记访问过的节点。
• 从目标启动 DFS：
  • 如果距离不是 K，则将值添加到最终总和中。
  • 如果该节点未被访问，则在par和每个节点的链接的帮助下，继续对其邻居（即父节点和子节点）进行 DFS 遍历。
  • 在当前节点的递归完成时返回其邻居的总和
• 返回距目标K距离内的所有节点的总和。

下面是上述方法的实现：

C++

// C++ code to implement above approach
  
#include 
using namespace std;
  
// Structure of a tree node
struct Node {
    int data;
    Node* left;
    Node* right;
    Node(int val)
    {
        this->data = val;
        this->left = 0;
        this->right = 0;
    }
};
  
// Function for marking the parent node
// for all the nodes using DFS
void dfs(Node* root,
         unordered_map& par)
{
    if (root == 0)
        return;
    if (root->left != 0)
        par[root->left] = root;
    if (root->right != 0)
        par[root->right] = root;
    dfs(root->left, par);
    dfs(root->right, par);
}
  
// Function calling for finding the sum
void dfs3(Node* root, int h, int& sum, int k,
          unordered_map& vis,
          unordered_map& par)
{
    if (h == k + 1)
        return;
    if (root == 0)
        return;
    if (vis[root])
        return;
    sum += root->data;
    vis[root] = 1;
    dfs3(root->left, h + 1, sum, k, vis, par);
    dfs3(root->right, h + 1, sum, k, vis, par);
    dfs3(par[root], h + 1, sum, k, vis, par);
}
  
// Function for finding
// the target node in the tree
Node* dfs2(Node* root, int target)
{
    if (root == 0)
        return 0;
    if (root->data == target)
        return root;
    Node* node1 = dfs2(root->left, target);
    Node* node2 = dfs2(root->right, target);
    if (node1 != 0)
        return node1;
    if (node2 != 0)
        return node2;
}
  
// Function to find the sum at distance K
int sum_at_distK(Node* root, int target,
                 int k)
{
    // Hash Table to store
    // the parent of a node
    unordered_map par;
  
    // Make the parent of root node as NULL
    // since it does not have any parent
    par[root] = 0;
  
    // Mark the parent node for all the
    // nodes using DFS
    dfs(root, par);
  
    // Find the target node in the tree
    Node* node = dfs2(root, target);
  
    // Hash Table to mark
    // the visited nodes
    unordered_map vis;
  
    int sum = 0;
  
    // DFS call to find the sum
    dfs3(node, 0, sum, k, vis, par);
    return sum;
}
  
// Driver Code
int main()
{
    // Taking Input
    Node* root = new Node(1);
    root->left = new Node(2);
    root->right = new Node(9);
    root->left->left = new Node(4);
    root->right->left = new Node(5);
    root->right->right = new Node(7);
    root->left->left->left = new Node(8);
    root->left->left->right = new Node(19);
    root->right->right->left = new Node(20);
    root->right->right->right
        = new Node(11);
    root->left->left->left->left
        = new Node(30);
    root->left->left->right->left
        = new Node(40);
    root->left->left->right->right
        = new Node(50);
  
    int target = 9, K = 1;
  
    // Function call
    cout << sum_at_distK(root, target, K);
    return 0;
}

输出

22

时间复杂度： O(N) 其中 N 是树中的节点数
辅助空间： O(N)