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> graph[max_n]; //与グラフの情報
vector<ll> getvs; //各DFSで見た頂点のリスト
ll answer[max_n]; //各色についての答え(の2倍)
ll tsize[max_n]; //各頂点を根とする部分木のサイズ
ll bc[max_n];
ll getsize[max_n]; //(重心を求める際に使用)
bool dead[max_n]; //各段階で見るべき頂点かどうか
bool cused[max_n];
bool uc[max_n];
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(!uc[color[now]]) {
    	bc[color[now]] += tsize[now];
      	uc[color[now]] = true;
    } else upper[now] = true;
  
    for(auto to: graph[now]) if(to != prev && !dead[to]) {
        dfs2(to, now);
    }
  
    if(!upper[now]) uc[color[now]] = false;
    else upper[now] = false;
}

//前計算その2 (その1と同様)
void dfs3(ll now, ll prev) {
    getvs.emplace_back(now);
    cused[color[now]] = true;
  
    if(!uc[color[now]]) {
      	bc[color[now]] += tsize[now];
      	uc[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]) uc[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);
        bc[color[centroid]]--;
      
        for(auto to: getvs) {
            if(us[color[to]]) {
            	answer[color[to]] -= bc[color[to]] * bc[color[to]];
              	//リセット
            	bc[color[to]] = 0;
            	cused[color[to]] = false;
            }
        }
      	//リセット
        getvs.clear();
      
      	//無関係な範囲を探索しないようにする
        dead[centroid] = true;
      
        for(auto to1: graph[centroid]) if(!dead[to1]) {
          	//重心を例外処理
          	cused[color[centroid]] = true;  
          	uc[color[centroid]] = true;
            bc[color[centroid]] = tsize[to1];
            getvs.emplace_back(centroid);
          
          	//前計算その2 (その1と同様)
            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]] += bc[color[to2]] * bc[color[to2]];
                  	//リセット
                    bc[color[to2]] = 0;
                    cused[color[to2]] = false;
                }
            }
          	//リセット
            uc[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:116:16: error: 'us' was not declared in this scope; did you mean 'uc'?
  116 |             if(us[color[to]]) {
      |                ^~
      |                uc