SumTree - Editorial

Practice :
https://www.codechef.com/problems/CODEPHR4

Contest Page :
https://www.codechef.com/COPR2020/problems/CODEPHR4

Author :
CodeAsylums

Tester :
CodeAsylums

Editorialist :
jalakp19

DIFFICULTY:
Simple

PREREQUISITES:
dfs, dp on graphs

PROBLEM:
Given a weighted tree of N vertices rooted at vertex 1, Find the sum of costs to reach each vertex from the root.

CONSTRAINTS:
1 <= N <= 10^5
1 <= u,v <= N
1 <= w <= 10^9

QUICK EXPLANATION:
Run dfs from the root node and calculate the size of each node(Number of nodes in the sub-tree of each node) and find the sum simultaneously.

EXPLANATION:
* KEY OBSERVATION *
Each edge weight is added to the total sum for every node in its subtree.
The Number of nodes in the sub-tree of each node can be calculated in a single dfs(Lets denote this by size of each node sz[]).
Now, we can calculate the Sum in the same dfs by adding the (edge-weight * size of the node) for each node.
Refer to the Solution below for better understanding.

SOLUTION:

//created by jalakp19

#include<bits/stdc++.h>
using namespace std;

typedef long long int ll;

#define sc(n) scanf("%lld",&n)
#define scc(n,m) scanf("%lld %lld",&n,&m)
#define sccc(x,y,z) scanf("%lld %lld %lld",&x,&y,&z)

vector mp[100005];
bool visit[100005];
ll sz[100005];
unordered_map<ll,unordered_map<ll,ll>> wt{};
ll sum{};

void dfs(ll u)
{
visit[u]=1;
sz[u]=1;
for(auto v:mp[u])
{
if(visit[v]==0)
{
dfs(v);
sz[u]+=sz[v];
sum+=(sz[v]*wt[u][v]);
}
}
}

int main()
{
ll n{};
sc(n);
for(ll i=0;i<n-1;i++)
{
ll x{},y{},w{};
sccc(x,y,w);
mp[x].push_back(y);
mp[y].push_back(x);
wt[x][y]=w;
wt[y][x]=w;
}
dfs(1);
cout<<sum<<endl;
return 0;
}