Caminhos.cpp

#include <bits/stdc++.h>

using namespace std;

const int maxn = 2e5+10;

typedef long long ll;

int n, c[maxn], sz[maxn], used[maxn];

ll ans[maxn];

vector<int> grafo[maxn];

void dfs(int u, int p = 0){
    sz[u] = 1;
    for(auto v : grafo[u]){
        if(v == p) continue;
        dfs(v, u);
        sz[u] += sz[v];
    }
}

void solve(int u, int p = 0){
    ans[c[u]] += 1LL*(n - sz[u] - used[c[u]] + 1) * sz[u];

    ll sum = sz[u] - 1;

    for(auto v : grafo[u]){
        if(v == p) continue;
        sum -= sz[v];
        ans[c[u]] += sz[v] * sum;
    }

    int previous_used = used[c[u]];

    for(auto v : grafo[u]){
        if(v == p) continue;
        used[c[u]] = n - sz[v];
        solve(v, u);
    }

    used[c[u]] = previous_used + sz[u];
}

int main(){
    ios_base::sync_with_stdio(false);cin.tie(NULL);cout.tie(NULL);

    cin >> n;

    for(int i=1; i<=n; i++) cin >> c[i];

    for(int i=1; i<n; i++){
        int u, v;

        cin >> u >> v;

        grafo[u].push_back(v);
        grafo[v].push_back(u);
    }

    dfs(1);
    solve(1);

    for(int i=1; i<=n; i++) cout << ans[i] << "\n";

    return 0;
}