PROBLEM LINK:
DIFFICULTY:
EASY
PREREQUISITES:
Depth-First-Search
PROBLEM:
Alice is playing a game on an undirected tree, the tree is made up of N nodes numbered from 1 to N. Each node has some coins kept on it, Alice wants to move all the coins on all nodes to a single node on the tree, to do this Alice can do the following move many times : pick any 1 coin on a node and move that coin to its adjacent node. Alice wants to calculate the minimum number of moves needed to move all coins to node i for all 1 \leq i \leq N. The number of coins on node i is A_i. Let f_i be the minimum number of moves to move all the coins on the tree from given starting position to node i. Then calculate f_i for 1 \leq i \leq N.
EXPLANATION:
The minimum moves are done when in each move any coin which is being moved must be moved to an adjacent node which is closer to our target node. Thus it is easy to see that
f_i= \sum_{k = 1}^{N} A_k *d(k,i)
where d(k,i) is the distance between the nodes k and i
Brute-force solution is to calculate each f_i for all i using DFS. But this will require N transversals 1 for each node and hence it will become quadratic. But we can do it in only 2 transversals by storing some information in our first transversal.
2 separate DFS transversals are needed.
Root the tree on any node.
In first transversal , we will create an array arr which stores, the total number of coins in the subtree of each node .
We will also calculate f_{root} in this transversal.
After this information is available after first transversal, another transversal
gives the value of all f_i by using the information stored in the above array arr,
by using the following property, suppose node i is parent of root j, then
f_j =(f_ i -arr[ j ] ) + (arr[ root ] - arr[ j ]).
TIME COMPLEXITY:
O(N)
SOLUTIONS:
Setter's Solution
#pragma GCC optimize "trapv"
#include<bits/stdc++.h>
using namespace std;
#define pb push_back
#define mp make_pair
#define vi vector<int>
#define vl vector<ll>
#define ss second
#define ff first
#define ll long long
#define pii pair<int,int>
#define pll pair<ll,ll>
#define tri pair<int,pii>
#define vii vector<pii>
#define vll vector<pll>
#define ld long double
#define mod 1000000007
int * coins;
ll * arr;
ll * f;
bool * visited;
vector<vl> adj;
int dfs(int i,int len=0)
{
visited[i]=true;
arr[i]=coins[i];
ll sum=0;
if(i==1 && adj[1].size()==1)
{
sum=dfs(adj[1][0],1);
arr[i]+=arr[adj[1][0]];
return sum;
}
if(adj[i].size()==1)
{
return 0;
}
for(int j=0;j<adj[i].size();j++)
{
if(!visited[adj[i][j]])
{
sum+=dfs(adj[i][j]);
arr[i]+=arr[adj[i][j]];
sum+=arr[adj[i][j]];
}
}
return sum;
}
void dfs2(int i)
{
visited[i]=true;
if(i==1 && adj[i].size()==1)
{
f[adj[1][0]]=f[1]+arr[1]-2*arr[adj[1][0]];
dfs2(adj[1][0]);
}
if(adj[i].size()==1)
return;
for(int j=0;j<adj[j].size();j++)
{
if(!visited[adj[i][j]])
{
f[adj[i][j]]=f[i]+arr[i]-2*arr[adj[1][0]];
dfs2(adj[i][j]);
}
}
return;
}
int main()
{
int t; cin>>t;
while(t--)
{
int n;
cin>>n;
coins=new int [n+1];
for(int i=0;i<n;i++)
cin>>coins[i+1];
if(n==1)
{
cout<<0<<endl;continue;
}
arr=new ll [n+1];
visited=new bool[n+1];
for(int i=0;i<=n;i++)
{
arr[i]=0;
visited[i]=false;
}
vl temp;
adj.clear();
for(int i=0;i<n+1;i++)
{
adj.pb(temp);
}
for(int i=0;i<n-1;i++){int a,b; cin>>a>>b;adj[a].pb(b);adj[b].pb(a);}
ll f_1=dfs(1);
f=new ll [n+1];
for(int i=0;i<=n;i++)
{
visited[i]=false;
f[i]=0;
}
f[1]=f_1;
dfs2(1);
for(int i=1;i<n;i++)
cout<<f[i]<<" ";
cout<<f[n]<<endl;
delete [] coins;
delete [] arr;
delete [] f;
delete [] visited;
}
return 0;
}