Language: C

#include <bits/stdc++.h>
using namespace std;
#pragma GCC target("avx2")
#pragma GCC optimize("O3")
#pragma GCC optimize("unroll-loops")
typedef long long ll;
typedef pair<ll,ll> pr;
typedef vector<ll> vc;
typedef unordered_map<ll,ll> umap;
#define pb emplace_back
#define rep(i,a,b) for(ll i=a;i<=b;i++)
#define repr(i,a,b) for(ll i=a;i>=b;i--)
#define reps(i,v) for(ll i=0;i<v.size();i++)
template<typename T> void chmin(T &a, const T &b) { a = min(a, b); }
template<typename T> void chmax(T &a, const T &b) { a = max(a, b); }
const ll max_n = 250005;
ll n, c, cr[max_n], full;
ll sz[max_n], depth[max_n], ans[max_n];
ll bg[max_n], bc[max_n];
bool dead[max_n], uc[max_n], us[max_n], ue[max_n];
vector<ll> g[max_n], gt, gs[max_n];

void dfs1(ll now, ll pr){
    bg[now] = 1;
    for(auto to: g[now]) if(to != pr && !dead[to]) {
        dfs1(to, now);
        bg[now] += bg[to];
    }
}     

void dfs2(ll now, ll pr){
    gt.pb(now);
    us[cr[now]] = true;
    if(!uc[cr[now]]) bc[cr[now]] += bg[now];
    else ue[now] = true;
    uc[cr[now]] = true;
    for(auto to: g[now]) if(to != pr && !dead[to]) {
        dfs2(to, now);
    }
    if(!ue[now]) uc[cr[now]] = false;
    ue[now] = false;
}

void dfs3(ll now, ll pr){
    gt.pb(now);
    us[cr[now]] = true;
    if(!uc[cr[now]]) bc[cr[now]] += bg[now];
    else ue[now] = true;
    uc[cr[now]] = true;
    for(auto to: g[now]) if(to != pr && !dead[to]) {
        dfs3(to, now);
    }
    if(cr[now] != cr[c] && !ue[now]) uc[cr[now]] = false;
    ue[now] = false;
}

int main(void) {
    cin >> n;
    rep(i,0,n-1){
        cin >> cr[i];  cr[i]--;
    }
    rep(i,1,n-1){
        ll aa, bb;
        cin >> aa >> bb;
        aa--; bb--;
        g[aa].pb(bb);  g[bb].pb(aa);
    }
    auto find_centroid = [&](ll root) {
        auto get_size = [&](auto &&self, ll v, ll p) -> void {
            sz[v] = 1;
            for(auto to: g[v]) if(to != p && !dead[to]) {
                self(self, to, v);
                sz[v] += sz[to];
            }
        };
        get_size(get_size, root, -1);
        auto dfs = [&](auto &&self, ll v, ll p) -> ll {
            for(auto to: g[v]) if(to != p && !dead[to]) {
                if(sz[to] > sz[root]/2) return self(self, to, v);
            }
            return v;
        };
        return dfs(dfs, root, -1);
    };
    auto centroid_decomposition = [&](auto &&self, ll root, ll d) -> void {
        c = find_centroid(root);
        dfs1(c, -1);
        //cout << "center " << c << " " << bg[c] << endl;
        dfs2(c, -1);
        bc[cr[c]]--;
        for(auto to: gt){
            if(us[cr[to]]) ans[cr[to]] -= bc[cr[to]] * bc[cr[to]];
            //if(us[cr[to]]) cout << " ans1 " << cr[to] << " " << ans[cr[to]] << endl;
            bc[cr[to]] = 0;
            us[cr[to]] = false;
        }
        gt.clear();
        dead[c] = true;
        depth[c] = d;
        for(auto to: g[c]) if(!dead[to]) {
            uc[cr[c]] = true;
            us[cr[c]] = true;
            bc[cr[c]] = bg[to];
            gt.pb(c);
            dfs3(to, c);
            for(auto to2: gt){
                if(us[cr[to2]]) ans[cr[to2]] += 2 * bc[cr[to2]] * (bg[c] - bg[to]);
                if(us[cr[to2]]) ans[cr[to2]] += bc[cr[to2]] * bc[cr[to2]];
                //if(us[cr[to2]]) cout << " ans2 " << cr[to2] << " " << ans[cr[to2]] << endl;
                bc[cr[to2]] = 0;
                us[cr[to2]] = false;
            }
            uc[cr[c]] = false;
            gt.clear();
        }
        ans[cr[c]] += 2;
        //cout << " ans3 " << cr[c] << " " << ans[cr[c]] << endl;
        for(auto to: g[c]) if(!dead[to]) {
            self(self, to, d+1);
        }
        dead[c] = false;
    };
    centroid_decomposition(centroid_decomposition, 0, 0);
    rep(i,0,n-1) cout << ans[i] / 2 << endl;
    return 0;
    }

Compiler output:
main.c:1:10: fatal error: bits/stdc++.h: No such file or directory
    1 | #include <bits/stdc++.h>
      |          ^~~~~~~~~~~~~~~
compilation terminated.