Language: C++

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
#define rep(i,a,b) for(ll i=a;i<=b;i++)

const ll max_n = 200005; //頂点をindexにもつ配列
const ll colors = 200005; //色をindexにもつ配列
ll n; //頂点数
ll centroid; //(ある時点での部分木の)重心
ll color[max_n]; //各頂点の色
vector<ll> graph[max_n]; //与グラフの情報
vector<ll> getvs; //各DFSで見た頂点のリスト
ll answer[colors]; //各色についての答え(の2倍)
ll tsize[max_n]; //各頂点を根とする部分木のサイズ
ll vsum[colors]; //ある色で最も根側にある頂点たちを根とする部分木のサイズの総和
ll getsize[max_n]; //(重心を求める際に使用)
bool dead[max_n]; //ある段階で見るべき頂点かどうか
bool cused[colors]; //ある色が探索した部分木全体で使われているか
bool upused[colors];  //根までさかのぼった時にある色が現れるか
bool upper[max_n]; //根までさかのぼった時に同色が現れないか

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

//前計算その1
void dfs2(ll now, ll prev) {
  	//通った頂点のリスト
    getvs.emplace_back(now);
  	//使われている色を管理
    cused[color[now]] = true;
  
    //根までさかのぼったときに同色が現れるか？
  	if(!upused[color[now]]) {
      	//dfs1で前計算したサイズを累積
    	vsum[color[now]] += tsize[now];
      	upused[color[now]] = true;
    } else upper[now] = true; //同色内でもっとも根側
  
    for(auto to: graph[now]) if(to != prev && !dead[to]) {
        dfs2(to, now);
    }
  
    if(!upper[now]) upused[color[now]] = false;
    else upper[now] = false;
}

//前計算その2 (その1と同様)
void dfs3(ll now, ll prev) {
    getvs.emplace_back(now);
    cused[color[now]] = true;
  
    if(!upused[color[now]]) {
      	vsum[color[now]] += tsize[now];
      	upused[color[now]] = true;
    } else upper[now] = true;
  
    for(auto to: graph[now]) if(to != prev && !dead[to]) {
        dfs3(to, now);
    }
  
    if(color[now] != color[centroid] && !upper[now]) upused[color[now]] = false;
    else upper[now] = false;
}

//ここからmain関数
int main(void) {
  
    //入力パート
    cin >> n;
    rep(i,1,n) cin >> color[i];   
    rep(i,1,n-1) {
        ll ainp, binp;
        cin >> ainp >> binp;
        graph[ainp].emplace_back(binp);
      	graph[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: graph[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: graph[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);
      
      	//前計算その1
        dfs2(centroid, 0);
      	//重心を端点とするパスの扱いに注意
        vsum[color[centroid]]--;
      
		//通った色についてのみ更新すれば良い      
        for(auto to: getvs) {
            if(cused[color[to]]) {
            	answer[color[to]] -= vsum[color[to]] * vsum[color[to]];
              	//リセット
            	vsum[color[to]] = 0;
            	cused[color[to]] = false;
            }
        }
      	//リセット
        getvs.clear();
      
      	//無関係な範囲を探索しないようにする
        dead[centroid] = true;
      
        for(auto to1: graph[centroid]) if(!dead[to1]) {
          	//重心を例外処理
          	cused[color[centroid]] = true;  
          	upused[color[centroid]] = true;
            vsum[color[centroid]] = tsize[to1];
            getvs.emplace_back(centroid);
          
          	//前計算その2 (その1と同様)
          	//ただしdfs2とは異なり, 重心の各子に対して実行
            dfs3(to1, centroid);
          	
            for(auto to2: getvs) {
                if(cused[color[to2]]) {
                    answer[color[to2]] += 2 * bc[color[to2]] * (tsize[centroid] - tsize[to1]);
				    answer[color[to2]] += vsum[color[to2]] * vsum[color[to2]];
                  	//リセット
                    vsum[color[to2]] = 0;
                    cused[color[to2]] = false;
                }
            }
          	//リセット
            upused[color[centroid]] = false;
            getvs.clear();
        }
      
      	//重心だけからなるパスを追加
        answer[color[centroid]] += 2;
      
      	//分割後の部分木でさらに重心分解
        for(auto to: graph[centroid]) if(!dead[to]) {
            self(self, to);
        }
      
      	//範囲の制限を解除
        dead[centroid] = false;
    };
  	//重心分解 ここまで
  
  	//適当な頂点(ここでは1)から重心分解を始める
    centroid_decomposition(centroid_decomposition, 1);
  
    //出力パート (半分にするのを忘れずに)
    rep(i,1,n) cout << answer[i] / 2 << endl;
  
    return 0;
}


Compiler output:
a.cc: In lambda function:
a.cc:147:47: error: 'bc' was not declared in this scope
  147 |                     answer[color[to2]] += 2 * bc[color[to2]] * (tsize[centroid] - tsize[to1]);
      |                                               ^~