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;
#define rep(i,a,b) for(ll i=a;i<=b;i++)
const ll max_n = 200005;
ll n; //頂点数
ll centroid; //(ある時点での部分木の)重心
ll color[max_n]; //各頂点の色
vector<ll> g[max_n]; //与グラフの情報
vector<ll> getvs; //各DFSで見た部分木に含まれる頂点のリスト
ll answer[max_n]; //各色についての答え(の2倍)
ll bg[max_n];
ll bc[max_n];
ll getsize[max_n]; //(重心を求める際に使用)
bool dead[max_n]; //各段階で見るべき頂点かどうか
bool uc[max_n];
bool us[max_n];
bool ue[max_n];

//各頂点を根とする部分木のサイズを求める
void dfs1(ll now, ll prev){
    bg[now] = 1;
    for(auto to: g[now]) if(to != prev && !dead[to]) {
        dfs1(to, now);
      	//各子のサイズを足す
        bg[now] += bg[to];
    }
}     

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

void dfs3(ll now, ll prev){
    getvs.emplace_back(now);
    us[color[now]] = true;
    if(!uc[color[now]]) bc[color[now]] += bg[now];
    else ue[now] = true;
    uc[color[now]] = true;
    for(auto to: g[now]) if(to != prev && !dead[to]) {
        dfs3(to, now);
    }
    if(color[now] != color[centroid] && !ue[now]) uc[color[now]] = false;
    ue[now] = false;
}

int main(void) {
    //入力パート
    cin >> n;
    rep(i,1,n) cin >> color[i];   
    rep(i,1,n-1){
        ll ainp, binp;
        cin >> ainp >> binp;
        g[ainp].emplace_back(binp);
      	g[binp].emplace_back(ainp);
    }
    //重心を求める (中の挙動は気にしなくて良い)
    auto find_centroid = [&](ll root) {
        auto get_size = [&](auto &&self, ll now, ll prev) -> void {
            getsize[now] = 1;
            for(auto to: g[now]) if(to != prev && !dead[to]) {
                self(self, to, now);
                getsize[now] += getsize[to];
            }
        };
        get_size(get_size, root, -1);
        auto dfs0 = [&](auto &&self, ll now, ll prev) -> ll {
            for(auto to: g[now]) if(to != prev && !dead[to]) {
                if(getsize[to] > getsize[root]/2) return self(self, to, now);
            }
            return now;
        };
        return dfs0(dfs0, root, 0);
    };
    //重心分解
    auto centroid_decomposition = [&](auto &&self, ll root) -> void {
        //重心を求める
        centroid = find_centroid(root);
      	//各頂点を根とする部分木のサイズを求める
        dfs1(centroid, 0);
        dfs2(centroid, 0);
        bc[color[centroid]]--;
        for(auto to: getvs){
            if(us[color[to]]) answer[color[to]] -= bc[color[to]] * bc[color[to]];
            bc[color[to]] = 0;
            us[color[to]] = false;
        }
        getvs.clear();
        dead[centroid] = true;
        for(auto to1: g[centroid]) if(!dead[to1]) {
            uc[color[centroid]] = true;
            us[color[centroid]] = true;
            bc[color[centroid]] = bg[to];
            getvs.emplace_back(centroid);
            dfs3(to1, centroid);
            for(auto to2: getvs){
                if(us[color[to2]]) answer[color[to2]] += 2 * bc[color[to2]] * (bg[centroid] - bg[to1]);
                if(us[color[to2]]) answer[color[to2]] += bc[color[to2]] * bc[color[to2]];
                bc[color[to2]] = 0;
                us[color[to2]] = false;
            }
            uc[color[centroid]] = false;
            getvs.clear();
        }
        ans[color[centroid]] += 2;
        for(auto to: g[centroid]) if(!dead[to]) {
            self(self, to);
        }
        dead[centroid] = false;
    };
    centroid_decomposition(centroid_decomposition, 1);
    //出力パート
    rep(i,1,n) cout << answer[i] / 2 << endl;
    return 0;
}


Compiler output:
a.cc: In function 'void dfs2(ll, ll)':
a.cc:36:28: error: 'cr' was not declared in this scope
   36 |     if(!uc[color[now]]) bc[cr[now]] += bg[now];
      |                            ^~
a.cc: In lambda function:
a.cc:105:38: error: 'to' was not declared in this scope; did you mean 'tm'?
  105 |             bc[color[centroid]] = bg[to];
      |                                      ^~
      |                                      tm
a.cc:117:9: error: 'ans' was not declared in this scope; did you mean 'abs'?
  117 |         ans[color[centroid]] += 2;
      |         ^~~
      |         abs