HackerEarth Sum of Sums problem solution

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.

#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;
}