查询以查找给定节点到加权树中每个叶节点的距离总和

给定具有N 个节点和E 条边的无向加权树。给定Q 个查询，每个查询指示一个起始节点。任务是打印从给定起始节点S到加权树中每个叶节点的距离总和。
例子：

Input:  N = 5, E = 4, Q = 3                    
                    1
            (4) /    \ (2)
              /        \
            4          2
                 (5)/    \ (3)
                  /        \
               5           3
Query 1: S = 1 
Query 2: S = 3 
Query 3: S = 5
Output : 
16
17
19
Explanation : 
The three leaf nodes in the tree are 3, 4 and 5. 
For S = 1, the sum of the distances from node 1 to the leaf nodes are: d(1, 4) + d(1, 3) + d(1, 5) = 4 + (2 + 5) + (2 + 3) = 16.
For S = 3, the sum of the distances from node 3 to its leaf nodes are: d(3, 4) + d(3, 3) + d(3, 5) = (3 + 2 + 4) + 0 + (3 + 5) = 17
For S = 5, the sum of the distances from node 5 to its leaf nodes are: d(5, 4) + d(5, 3) + d(5, 5) = (5 + 2 + 4) + (5 + 3) + 0 = 19

Input:  N = 3, E = 2, Q = 2                    
                    1             
            (9) /    \ (1)
              /        \  
           2           3  
Query 1: S = 1 
Query 2: S = 2 
Query 3: S = 3 
Output : 
10
10
10 

天真的方法：
对于每个查询，遍历整个树并找到从给定源节点到所有叶节点的距离总和。

时间复杂度： O(Q * N)

高效方法：思想是使用动态规划树算法预先计算每个节点到所有叶节点的距离之和，并在恒定时间内获得每个查询的答案。

请按照以下步骤解决问题

• 初始化一个向量dp来存储从每个节点i到树的所有叶节点的距离之和。
• 初始化一个向量叶子以存储节点i的子树中的叶子节点数，以1为根节点。
• 查找节点i的距离的总和在我使用改进的深度优先搜索算法考虑1为根节点的子树中的所有叶子节点。

Let node a be the parent of node i

• leaves[a] += leaves[i] ;
• dp[a] += dp[i] + leaves[i] * weight of edge between nodes(a, i) ;

• 使用 re-rooting 技术找到不在节点i的子树中的树的剩余叶子的距离。要计算这些距离，请使用另一种修改后的深度优先搜索 (DFS) 算法来查找叶节点到节点i的距离之和并将其相加。

Let a be the parent node and i be the child node, then
Let the number of leaf nodes outside the sub-tree i that are present in the sub-tree a be L

• L = leaves[a] – leaves[i] ;
• dp[i] += ( dp[a] – dp[i] ) + ( weight of edge between nodes(a, i) ) * ( L – leaves[i] ) ;
• leaves[i] += L ;

下面是上述方法的实现：

C++

// C++ program for the above problem
#include 
using namespace std;
 
// MAX size
const int N = 1e5 + 5;
 
// graph with {destination, weight};
vector > > v(N);
 
// for storing the sum for ith node
vector dp(N);
 
// leaves in subtree of ith.
vector leaves(N);
int n;
 
// dfs to find sum of distance
// of leaves in the
// subtree of a node
void dfs(int a, int par)
{
    // flag is the node is
    // leaf or not;
    bool leaf = 1;
    for (auto& i : v[a]) {
        // skipping if parent
        if (i.first == par)
            continue;
 
        // setting flag to false
        leaf = 0;
 
        // doing dfs call
        dfs(i.first, a);
    }
 
    // doing calculation
    // in postorder.
    if (leaf == 1) {
 
        // if the node is leaf then
        // we just increment
        // the no. of leaves under
        // the subtree of a node
        leaves[a] += 1;
    }
    else {
 
        for (auto& i : v[a]) {
            if (i.first == par)
                continue;
 
            // adding num of leaves
            leaves[a]
                += leaves[i.first];
 
            // calculating answer for
            // the sum in the subtree
            dp[a] = dp[a]
                    + dp[i.first]
                    + leaves[i.first]
                          * i.second;
        }
    }
}
 
// dfs function to find the
// sum of distance of leaves
// outside the subtree
void dfs2(int a, int par)
{
    for (auto& i : v[a]) {
        if (i.first == par)
            continue;
 
        // number of leaves other
        // than the leaves in the
        // subtree of i
        int leafOutside = leaves[a] - leaves[i.first];
 
        // adding the contribution
        // of leaves outside to
        // the ith node
        dp[i.first] += (dp[a] - dp[i.first]);
 
        dp[i.first] += i.second
                       * (leafOutside
                          - leaves[i.first]);
 
        // adding the leafs outside
        // to ith node's leaves.
        leaves[i.first]
            += leafOutside;
        dfs2(i.first, a);
    }
}
 
void answerQueries(
    vector queries)
{
 
    // calculating the sum of
    // distance of leaves in the
    // subtree of a node assuming
    // the root of the tree is 1
    dfs(1, 0);
 
    // calculating the sum of
    // distance of leaves outside
    // the subtree of node
    // assuming the root of the
    // tree is 1
    dfs2(1, 0);
 
    // answering the queries;
    for (int i = 0;
         i < queries.size(); i++) {
        cout << dp[queries[i]] << endl;
    }
}
 
// Driver Code
int main()
{
    // Driver Code
    /*
             1
       (4) /   \ (2)
          /     \
         4       2
             (5)/  \ (3)
               /    \
              5      3
    */
 
    n = 5;
 
    // initialising tree
    v[1].push_back(
        make_pair(4, 4));
    v[4].push_back(
        make_pair(1, 4));
    v[1].push_back(
        make_pair(2, 2));
    v[2].push_back(
        make_pair(1, 2));
    v[2].push_back(
        make_pair(3, 3));
    v[3].push_back(
        make_pair(2, 3));
    v[2].push_back(
        make_pair(5, 5));
    v[5].push_back(
        make_pair(2, 5));
 
    vector
        queries = { 1, 3, 5 };
    answerQueries(queries);
}

Java

// Java program for the above problem
import java.util.*;
import java.lang.*;
 
class GFG{
     
static class pair
{
    int first, second;
     
    pair(int f, int s)
    {
        this.first = f;
        this.second = s;
    }
}
 
// MAX size
static final int N = (int)1e5 + 5;
  
// Graph with {destination, weight};
static ArrayList> v;
  
// For storing the sum for ith node
static int[] dp = new int[N];
  
// Leaves in subtree of ith.
static int[] leaves = new int[N];
static int n;
  
// dfs to find sum of distance
// of leaves in the subtree of
// a node
static void dfs(int a, int par)
{
     
    // Flag is the node is
    // leaf or not;
    int leaf = 1;
     
    for(pair i : v.get(a))
    {
         
        // Skipping if parent
        if (i.first == par)
            continue;
  
        // Setting flag to false
        leaf = 0;
  
        // Doing dfs call
        dfs(i.first, a);
    }
  
    // Doing calculation
    // in postorder.
    if (leaf == 1)
    {
         
        // If the node is leaf then
        // we just increment the
        // no. of leaves under
        // the subtree of a node
        leaves[a] += 1;
    }
    else
    {
        for(pair i : v.get(a))
        {
            if (i.first == par)
                continue;
  
            // Adding num of leaves
            leaves[a] += leaves[i.first];
  
            // Calculating answer for
            // the sum in the subtree
            dp[a] = dp[a] + dp[i.first] +
                        leaves[i.first] *
                               i.second;
        }
    }
}
  
// dfs function to find the
// sum of distance of leaves
// outside the subtree
static void dfs2(int a, int par)
{
    for(pair i : v.get(a))
    {
        if (i.first == par)
            continue;
  
        // Number of leaves other
        // than the leaves in the
        // subtree of i
        int leafOutside = leaves[a] -
                          leaves[i.first];
  
        // Adding the contribution
        // of leaves outside to
        // the ith node
        dp[i.first] += (dp[a] - dp[i.first]);
  
        dp[i.first] += i.second *
                   (leafOutside -
                    leaves[i.first]);
  
        // Adding the leafs outside
        // to ith node's leaves.
        leaves[i.first] += leafOutside;
        dfs2(i.first, a);
    }
}
  
static void answerQueries(int[] queries)
{
     
    // Calculating the sum of
    // distance of leaves in the
    // subtree of a node assuming
    // the root of the tree is 1
    dfs(1, 0);
  
    // Calculating the sum of
    // distance of leaves outside
    // the subtree of node
    // assuming the root of the
    // tree is 1
    dfs2(1, 0);
  
    // Answering the queries;
    for(int i = 0; i < queries.length; i++)
    {
        System.out.println(dp[queries[i]]);
    }
}
 
// Driver code
public static void main(String[] args)
{
     
    /*
             1
       (4) /   \ (2)
          /     \
         4       2
             (5)/  \ (3)
               /    \
              5      3
    */
     
    n = 5;
     
    v = new ArrayList<>();
    for(int i = 0; i <= n; i++)
        v.add(new ArrayList());
     
    // Initialising tree
    v.get(1).add(new pair(4, 4));
    v.get(4).add(new pair(1, 4));
    v.get(1).add(new pair(2, 2));
    v.get(2).add(new pair(1, 2));
    v.get(2).add(new pair(3, 3));
    v.get(3).add(new pair(2, 3));
    v.get(2).add(new pair(5, 5));
    v.get(5).add(new pair(2, 5));
     
    // System.out.println(v);
    int[] queries = { 1, 3, 5 };
     
    answerQueries(queries);
}
}
 
// This code is contributed by offbeat

Python3

# Python3 program for the above problem
 
# MAX size
N = 10 ** 5 + 5
 
# Graph with {destination, weight}
v = [[] for i in range(N)]
 
# For storing the sum for ith node
dp = [0] * N
 
# Leaves in subtree of ith.
leaves = [0] * (N)
n = 0
 
# dfs to find sum of distance
# of leaves in the subtree of
# a node
def dfs(a, par):
     
    # flag is the node is
    # leaf or not
    leaf = 1
     
    for i in v[a]:
         
        # Skipping if parent
        if (i[0] == par):
            continue
 
        # Setting flag to false
        leaf = 0
 
        # Doing dfs call
        dfs(i[0], a)
 
    # Doing calculation
    # in postorder.
    if (leaf == 1):
         
        # If the node is leaf then
        # we just increment
        # the no. of leaves under
        # the subtree of a node
        leaves[a] += 1
    else:
        for i in v[a]:
            if (i[0] == par):
                continue
 
            # Adding num of leaves
            leaves[a] += leaves[i[0]]
 
            # Calculating answer for
            # the sum in the subtree
            dp[a] = (dp[a] + dp[i[0]] +
                 leaves[i[0]] * i[1])
 
# dfs function to find the
# sum of distance of leaves
# outside the subtree
def dfs2(a, par):
     
    for i in v[a]:
        if (i[0] == par):
            continue
 
        # Number of leaves other
        # than the leaves in the
        # subtree of i
        leafOutside = leaves[a] - leaves[i[0]]
 
        # Adding the contribution
        # of leaves outside to
        # the ith node
        dp[i[0]] += (dp[a] - dp[i[0]])
 
        dp[i[0]] += i[1] * (leafOutside -
                            leaves[i[0]])
 
        # Adding the leafs outside
        # to ith node's leaves.
        leaves[i[0]] += leafOutside
        dfs2(i[0], a)
 
def answerQueries(queries):
 
    # Calculating the sum of
    # distance of leaves in the
    # subtree of a node assuming
    # the root of the tree is 1
    dfs(1, 0)
 
    # Calculating the sum of
    # distance of leaves outside
    # the subtree of node
    # assuming the root of the
    # tree is 1
    dfs2(1, 0)
 
    # Answering the queries
    for i in range(len(queries)):
        print(dp[queries[i]])
 
# Driver Code
if __name__ == '__main__':
     
    #          1
    #    (4) /   \ (2)
    #       /     \
    #      4       2
    #          (5)/  \ (3)
    #            /    \
    #           5      3
    #
 
    n = 5
 
    # Initialising tree
    v[1].append([4, 4])
    v[4].append([1, 4])
    v[1].append([2, 2])
    v[2].append([1, 2])
    v[2].append([3, 3])
    v[3].append([2, 3])
    v[2].append([5, 5])
    v[5].append([2, 5])
 
    queries = [ 1, 3, 5 ]
     
    answerQueries(queries)
 
# This code is contributed by mohit kumar 29

Javascript

 
// JavaScript program for the above problem
 
class pair
{
    constructor(f,s)
    {
        this.first = f;
        this.second = s;
    }
}
 
// MAX size
let N = 1e5 + 5;
 
// Graph with {destination, weight};
let v=[];
 
// For storing the sum for ith node
let dp = new Array(N);
 
 
// Leaves in subtree of ith.
let leaves = new Array(N);
for(let i=0;i<N;i++)
{
    dp[i]=0;
    leaves[i]=0;
}
 
let n;
 
// dfs to find sum of distance
// of leaves in the subtree of
// a node
function dfs(a,par)
{
    // Flag is the node is
    // leaf or not;
    let leaf = 1;
      
    for(let i=0;i<v[a].length;i++)
    {
          
        // Skipping if parent
        if (v[a][i].first == par)
            continue;
   
        // Setting flag to false
        leaf = 0;
   
        // Doing dfs call
        dfs(v[a][i].first, a);
    }
   
    // Doing calculation
    // in postorder.
    if (leaf == 1)
    {
          
        // If the node is leaf then
        // we just increment the
        // no. of leaves under
        // the subtree of a node
        leaves[a] += 1;
    }
    else
    {
        for(let i=0;i<v[a].length;i++)
        {
            if (v[a][i].first == par)
                continue;
   
            // Adding num of leaves
            leaves[a] += leaves[v[a][i].first];
   
            // Calculating answer for
            // the sum in the subtree
            dp[a] = dp[a] + dp[v[a][i].first] +
                        leaves[v[a][i].first] *
                               v[a][i].second;
        }
    }
}
 
// dfs function to find the
// sum of distance of leaves
// outside the subtree
function dfs2(a,par)
{
    for(let i=0;i< v[a].length;i++)
    {
        if (v[a][i].first == par)
            continue;
   
        // Number of leaves other
        // than the leaves in the
        // subtree of i
        let leafOutside = leaves[a] -
                          leaves[v[a][i].first];
   
        // Adding the contribution
        // of leaves outside to
        // the ith node
        dp[v[a][i].first] += (dp[a] - dp[v[a][i].first]);
   
        dp[v[a][i].first] += v[a][i].second *
                   (leafOutside -
                    leaves[v[a][i].first]);
   
        // Adding the leafs outside
        // to ith node's leaves.
        leaves[v[a][i].first] += leafOutside;
        dfs2(v[a][i].first, a);
    }
}
 
function answerQueries(queries)
{
    // Calculating the sum of
    // distance of leaves in the
    // subtree of a node assuming
    // the root of the tree is 1
    dfs(1, 0);
   
    // Calculating the sum of
    // distance of leaves outside
    // the subtree of node
    // assuming the root of the
    // tree is 1
    dfs2(1, 0);
   
    // Answering the queries;
    for(let i = 0; i < queries.length; i++)
    {
        document.write(dp[queries[i]]+"<br>");
    }
}
 
// Driver code
/*
             1
       (4) /   \ (2)
          /     \
         4       2
             (5)/  \ (3)
               /    \
              5      3
    */
 
n = 5;
 
v = [];
for(let i = 0; i <= n; i++)
    v.push([]);
 
// Initialising tree
v[1].push(new pair(4, 4));
v[4].push(new pair(1, 4));
v[1].push(new pair(2, 2));
v[2].push(new pair(1, 2));
v[2].push(new pair(3, 3));
v[3].push(new pair(2, 3));
v[2].push(new pair(5, 5));
v[5].push(new pair(2, 5));
 
// System.out.println(v);
let queries = [ 1, 3, 5 ];
 
answerQueries(queries);
 
 
// This code is contributed by patel2127
 

输出：

16
17
19

时间复杂度： O(N + Q)
辅助空间： O(N)

如果您希望与专家一起参加现场课程，请参阅DSA 现场工作专业课程和学生竞争性编程现场课程。