In this HackerEarth Sum of Sums problem solution Given a undirected tree containing N nodes where ith node has value a(i). Let us define cost of the tree as C, where
C = Summation(i=1, N) f(i)
f(i) = Summation(j) g(j); where j belongs to subtree of i
g(j) = Summation(k) a(k); where k belongs t subtree of j
Find a root of the tree such that the cost of the tree is minimum.


HackerEarth Sum of Sums problem solution


HackerEarth Sum of Sums problem solution.

#include<bits/stdc++.h>
using namespace std;
#define N 100009
long long a[N],g[N],h[N],ans,ans2,mx;
vector<int> G[N];
void dfs1(int u,int p){
h[u]=a[u];
g[u]=0;
for(int v:G[u]){
if(v==p)
continue;
dfs1(v,u);
h[u]+=h[v];
g[u]+=g[v];
}
g[u]+=h[u];
// cout<<"u : "<<u<<" "<<" , h[u] : "<<h[u]<<" , g[u] : "<<g[u]<<endl;
return ;
}

void dfs2(int u,int p){
//cout<<"u : "<<u<<" "<<" , ans : "<<ans<<endl;
if (mx > ans) {
mx = ans;
ans2 = u;
}
else if (mx == ans && ans2 > u)
ans2 = u;
for(int v:G[u]){
if(v==p)
continue;
g[u]-=g[v];ans-=g[v];
g[u]-=h[u];ans-=h[u];
h[u]-=h[v];
g[u]+=h[u];ans+=h[u];
g[v]+=g[u];ans+=g[u];
g[v]-=h[v];ans-=h[v];
h[v]+=h[u];
g[v]+=h[v];ans+=h[v];
dfs2(v,u);
g[v]-=h[v];ans-=h[v];
h[v]-=h[u];
g[v]+=h[v];ans+=h[v];
g[v]-=g[u];ans-=g[u];
g[u]-=h[u];ans-=h[u];
h[u]+=h[v];
g[u]+=h[u];ans+=h[u];
g[u]+=g[v];ans+=g[v];
}
return ;
}

int main(){
//freopen("in00.txt", "r", stdin);
//freopen("out14.txt", "w", stdout);
ios_base::sync_with_stdio(false);cin.tie(NULL);
int n;
cin>>n;
for(int i=0;i<n;i++){
cin>>a[i];
}
for(int i=0;i<n-1;i++){
int u,v;
cin>>u>>v;
u--;v--;
G[u].push_back(v);
G[v].push_back(u);
}
dfs1(0,-1);
ans=0;
for(int i=0;i<n;i++){
ans+=g[i];
}
mx=LLONG_MAX;
ans2=-1;
dfs2(0,-1);
cout<<ans2+1<<" "<<mx<<endl;
return 0;
}

Second solution

#include <string>
#include <vector>
#include <map>
#include <list>
#include <iterator>
#include <cassert>
#include <set>
#include <queue>
#include <iostream>
#include <sstream>
#include <stack>
#include <deque>
#include <cmath>
#include <memory.h>
#include <cstdlib>
#include <cstdio>
#include <cctype>
#include <algorithm>
#include <utility>
#include <time.h>
#include <complex>
using namespace std;

#define FOR(i, a, b) for(int i=(a);i<(b);i++)
#define RFOR(i, b, a) for(int i=(b)-1;i>=(a);--i)
#define FILL(A,value) memset(A,value,sizeof(A))

#define ALL(V) V.begin(), V.end()
#define SZ(V) (int)V.size()
#define PB push_back
#define MP make_pair
#define Pi 3.14159265358979
#define x0 ikjnrmthklmnt
#define y0 lkrjhkltr
#define y1 ewrgrg

typedef long long Int;
typedef unsigned long long UInt;
typedef vector<int> VI;
typedef pair<int, int> PII;
typedef pair<Int, Int> PLL;
typedef pair<double, double> PDD;
typedef complex<double> base;

const int INF = 1000000000;
const int BASE = 1000000007;
const int MAX = 100007;
const int ADD = 1000000;
const int MOD = 1000000007;
const int CNT = 800;

const double eps = 1e-8;

VI G[MAX];

Int f[MAX];
Int g[MAX];


int a[MAX];

void dfs(int v, int p)
{
g[v] = a[v];

FOR(i,0,SZ(G[v]))
{
int to = G[v][i];
if (to == p) continue;
dfs(to, v);
g[v] += g[to];
f[v] += f[to];
}
f[v] += g[v];
}


Int cur, M , id;

void dfs2(int v, int p)
{
//cout << v << ' ' << cur << endl;


if (cur < M)
{
M = cur;
id = v + 1;
}

FOR(i,0,SZ(G[v]))
{
int to = G[v][i];
if (to == p) continue;

Int cc = cur;
Int ffv = f[v];
Int fft = f[to];
Int ggv = g[v];
Int ggt = g[to];

cur -= f[v];
cur -= f[to];
f[v] -= f[to];
f[to] -= g[to];
f[v] -= g[v];
g[v] -= g[to];
f[v] += g[v];
f[to] += f[v];
g[to] += g[v];
f[to] += g[to];
cur += f[v];
cur += f[to];

//cout << f[v] << ' ' << g[v] << ' ' << f[to] << ' ' << g[to] << endl;

dfs2(to, v);

cur = cc;
f[v] = ffv;
f[to] = fft;
g[v] = ggv;
g[to] = ggt;

}

}

int main()
{
//freopen("in.txt", "r", stdin);
//freopen("distance.in", "r", stdin);
//freopen("distance.out", "w", stdout);
//freopen("out.txt" , "w" , stdout);

int n;
cin >> n;
assert(n >= 1 && n <= 100000);

FOR(i,0,n)
{
scanf("%d", &a[i]);
assert(a[i] >= 1 && a[i] <= 1000000);
}

FOR(i,0,n - 1)
{
int a , b;
scanf("%d%d" , &a , &b);
assert(a >= 1 && a <= n);
assert(b >= 1 && b <= n);
-- a;
-- b;
G[a].push_back(b);
G[b].push_back(a);
}


dfs(0,-1);
FOR(i,0,n)
{
assert(f[i] >= 1);
}


cur = 0;
FOR(i,0,n)
{
cur += f[i];
}

M = cur;
id = 1;

dfs2(0,-1);

cout << id << ' ' << M << endl;

return 0;
}