Cod sursa(job #2003526)

Utilizator MaligMamaliga cu smantana Malig Data 23 iulie 2017 09:05:19
Problema Paduri de multimi disjuncte Scor 0
Compilator cpp Status done
Runda Arhiva educationala Marime 1.63 kb
#include <iostream>
#include <fstream>
#include <cstring>
#include <vector>

#define ll long long
#define pb push_back
using namespace std;
ifstream in("sumOfAllDistances.in");
ofstream out("sumOfAllDistances.out");
const int NMax = 1e5 + 5;

int N;
ll sumSub[NMax],sumAll[NMax],nrSub[NMax];

struct elem {
    ll y,c;
    elem (ll _y = 0,ll _c = 0) {
        y = _y;
        c = _c;
    }
};
vector<elem> v[NMax];

void dfs1(int,int);
void dfs2(int,int,int);

int main() {
    in>>N;
    for (int i=1;i <= N-1;++i) {
        int x,y,c;
        in>>x>>y>>c;

        v[x].pb(elem(y,c));
        v[y].pb(elem(x,c));
    }

    dfs1(1,0);

    sumAll[1] = sumSub[1];
    dfs2(1,0,0);

    for (int i=1;i <= N;++i) {
        out<<sumAll[i]<<'\n';
    }

    in.close();out.close();
    return 0;
}

void dfs1(int node,int dad) {
    nrSub[node] = 1;
    sumSub[node] = 0;

    for (elem el : v[node]) {
        int son = el.y, cost = el.c;
        if (son == dad) {
            continue;
        }

        dfs1(son,node);
        nrSub[node] += nrSub[son];
        sumSub[node] += cost * nrSub[son] + sumSub[son];
    }
}

void dfs2(int node,int dad,int costBack) {
    if (node != 1) {
        sumAll[node] = sumSub[node] +
                       (
                        sumAll[dad] - (costBack * nrSub[node] + sumSub[node])
                        + (costBack * (N - nrSub[node]))
                        );
    }

    for (elem el : v[node]) {
        int son = el.y, cost = el.c;
        if (son == dad) {
            continue;
        }

        dfs2(son,node,cost);
    }
}