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> get; //各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){
    get.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){
    get.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 dfs(dfs, 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: get){
            if(us[color[to]]) answer[color[to]] -= bc[color[to]] * bc[color[to]];
            bc[color[to]] = 0;
            us[color[to]] = false;
        }
        get.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];
            get.emplace_back(centroid);
            dfs3(to1, centroid);
            for(auto to2: get){
                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;
            get.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 << ans[i] / 2 << endl;
    return 0;
}


Compiler output:
a.cc: In function 'void dfs2(ll, ll)':
a.cc:34:5: error: reference to 'get' is ambiguous
   34 |     get.emplace_back(now);
      |     ^~~
In file included from /usr/include/x86_64-linux-gnu/c++/14/bits/stdc++.h:80,
                 from a.cc:1:
/usr/include/c++/14/variant:1180:5: note: candidates are: 'template<class _Tp, class ... _Types> constexpr const _Tp&& std::get(const variant<_Types ...>&&)'
 1180 |     get(const variant<_Types...>&& __v)
      |     ^~~
/usr/include/c++/14/variant:1170:5: note:                 'template<class _Tp, class ... _Types> constexpr const _Tp& std::get(const variant<_Types ...>&)'
 1170 |     get(const variant<_Types...>& __v)
      |     ^~~
/usr/include/c++/14/variant:1160:5: note:                 'template<class _Tp, class ... _Types> constexpr _Tp&& std::get(variant<_Types ...>&&)'
 1160 |     get(variant<_Types...>&& __v)
      |     ^~~
/usr/include/c++/14/variant:1150:5: note:                 'template<class _Tp, class ... _Types> constexpr _Tp& std::get(variant<_Types ...>&)'
 1150 |     get(variant<_Types...>& __v)
      |     ^~~
/usr/include/c++/14/variant:1748:5: note:                 'template<long unsigned int _Np, class ... _Types> constexpr std::variant_alternative_t<_Np, std::variant<_Types ...> >&& std::get(const variant<_Types ...>&&)'
 1748 |     get(const variant<_Types...>&& __v)
      |     ^~~
/usr/include/c++/14/variant:1737:5: note:                 'template<long unsigned int _Np, class ... _Types> constexpr std::variant_alternative_t<_Np, std::variant<_Types ...> >& std::get(const variant<_Types ...>&)'
 1737 |     get(const variant<_Types...>& __v)
      |     ^~~
/usr/include/c++/14/variant:1726:5: note:                 'template<long unsigned int _Np, class ... _Types> constexpr std::variant_alternative_t<_Np, std::variant<_Types ...> >&& std::get(variant<_Types ...>&&)'
 1726 |     get(variant<_Types...>&& __v)
      |     ^~~
/usr/include/c++/14/variant:1715:5: note:                 'template<long unsigned int _Np, class ... _Types> constexpr std::variant_alternative_t<_Np, std::variant<_Types ...> >& std::get(variant<_Types ...>&)'
 1715 |     get(variant<_Types...>& __v)
      |     ^~~
In file included from /usr/include/c++/14/bits/memory_resource.h:47,
                 from /usr/include/c++/14/string:68,
                 from /usr/include/c++/14/bitset:52,
                 from /usr/include/x86_64-linux-gnu/c++/14/bits/stdc++.h:52:
/usr/include/c++/14/tuple:2517:5: note:                 'template<class _Tp, class ... _Types> constexpr const _Tp&& std::get(const tuple<_Elements ...>&&)'
 2517 |     get(const tuple<_Types...>&& __t) noexcept
      |     ^~~
/usr/include/c++/14/tuple:2505:5: note:                 'template<class _Tp, class ... _Types> constexpr const _Tp& std::get(const tuple<_Elements ...>&)'
 2505 |     get(const tuple<_Types...>& __t) noexcept
      |     ^~~
/usr/include/c++/14/tuple:2494:5: note:                 'template<class _Tp, class ... _Types> constexpr _Tp&& std::get(tuple<_Elements ...>&&)'
 2494 |     get(tuple<_Types...>&& __t) noexcept
      |     ^~~
/usr/include/c++/14/tuple:2483:5: note:                 'template<class _Tp, class ... _Types> constexpr _Tp& std::get(tuple<_Elements ...>&)'
 2483 |     get(tuple<_Types...>& __t) noexcept
      |     ^~~
/usr/include/c++/14/tuple:2476:5: note:                 'template<long unsigned int __i, class ... _Elements> constexpr std::__enable_if_t<(__i >= sizeof... (_Types))> std::get(const tuple<_Elements ...>&)'
 2476 |     get(const tuple<_Elements...>&) = delete;
      |     ^~~
In file included from /usr/include/c++/14/bits/stl_algobase.h:64,
                 from /usr/include/c++/14/algorithm:60,
                 from /usr/include/x86_64-linux-gnu/c++/14/bits/stdc++.h:51:
/usr/include/c++/14/bits/stl_pair.h:1307:5: note:                 'template<class _Tp, class _Up> constexpr const _Tp&& std::get(const pair<_Up, _Tp>&&)'
 1307 |     get(const pair<_Up, _Tp>&& __p) noexcept
      |     ^~~
/usr/include/c++/14/bits/stl_pair.h:1302:5: note:                 'template<class _Tp, class _Up> constexpr _Tp&& std::get(pair<_Up, _Tp>&&)'
 1302 |     get(pair<_Up, _Tp>&& __p) noexcept
      |     ^~~
/usr/include/c++/14/bits/stl_pair.h:1297:5: note:                 'template<class _Tp, class _Up> constexpr const _Tp& std::get(const pair<_Up, _Tp>&)'
 1297 |     get(const pair<_Up, _Tp>& __p) noexcept
      |     ^~~
/usr/include/c++/14/bits/stl_pair.h:1292:5: note:                 'template<class _Tp, class _Up> constexpr _Tp& std::get(pair<_Up, _Tp>&)'
 1292 |     get(pair<_Up, _Tp>& __p) noexcept
      |     ^~~
/usr/include/c++/14/bits/stl_pair.h:1287:5: note:                 'template<class _Tp, class _Up> constexpr const _Tp&& std::get(const pair<_T1, _T2>&&)'
 1287 |     get(const pair<_Tp, _Up>&& __p) noexcept
      |     ^~~
/usr/include/c++/14/bits/stl_pair.h:1282:5: note:                 'template<class _Tp, class _Up> constexpr _Tp&& std::get(pair<_T1, _T2>&&)'
 1282 |     get(pair<_Tp, _Up>&& __p) noexcept
      |     ^~~
/usr/include/c++/14/bits/stl_pair.h:1277:5: note:                 'template<class _Tp, class _Up> constexpr const _Tp& std::get(const pair<_T1, _T2>&)'
 1277 |     get(const pair<_Tp, _Up>& __p) noexcept
      |     ^~~
/usr/include/c++/14/bits/stl_pair.h:1272:5: note:                 'template<class _Tp, class _Up> constexpr _Tp& std::get(pair<_T1, _T2>&)'
 1272 |     get(pair<_Tp, _Up>& __p) noexcept
      |     ^~~
In file included from /usr/include/c++/14/functional:65,
                 from /usr/include/x86_64-linux-gnu/c++/14/bits/stdc++.h:53:
/usr/include/c++/14/array:417:5: note:                 'template<long unsigned int _Int, class _Tp, long unsigned int _Nm> constexpr const _Tp&& std::get(const array<_Tp, _Nm>&&)'
  417 |     get(const array<_Tp, _Nm>&& __arr) noexcept
      |     ^~~
/usr/include/c++/14/array:408:5: note:                 'template<long unsigned int _Int, class _Tp, long unsigned int _Nm> constexpr const _Tp& std::get(const array<_Tp, _Nm>&)'
  408 |     get(const array<_Tp, _Nm>& __arr) noexcept
      |     ^~~
/usr/include/c++/14/array:399:5: note:                 'template<long unsigned int _Int, class _Tp, long unsigned int _Nm> constexpr _Tp&& std::get(array<_Tp, _Nm>&&)'
  399 |     get(array<_Tp, _Nm>&& __arr) noexcept
      |     ^~~
/usr/include/c++/14/array:390:5: note:                 'template<long unsigned int _Int, class _Tp, long unsigned int _Nm> constexpr _Tp& std::get(array<_Tp, _Nm>&)'
  390 |     get(array<_Tp, _Nm>& __arr) noexcept
      |     ^~~
/usr/include/c++/14/tuple:2466:5: note:                 'template<long unsigned int __i, class ... _Elements> constexpr std::__tuple_element_t<__i, std::tuple<_Elements ...> >&& std::get(const tuple<_Elements ...>&&)'
 2466 |     get(const tuple<_Elements...>&& __t) noexcept
      |     ^~~
/usr/include/c++/14/tuple:2457:5: note:                 'template<long unsigned int __i, class ... _Elements> constexpr std::__tuple_element_t<__i, std::tuple<_Elements ...> >&& std::get(tuple<_Elements ...>&&)'
 2457 |     get(tuple<_Elements...>&& __t) noexcept
      |     ^~~
/usr/include/c++/14/tuple:2451:5: note:                 'template<long unsigned int __i, class ... _Elements> constexpr std::__tuple_element_t<__i, std::tuple<_Elements ...> >& std::get(const tuple<_Elements ...>&)'
 2451 |     get(const tuple<_Elements...>& __t) noexcept
      |     ^~~
/usr/include/c++/14/tuple:2445:5: note:                 'template<long unsigned int __i, class ... _Elements> constexpr std::__tuple_element_t<__i, std::tuple<_Elements ...> >& std::get(tuple<_Elements ...>&)'
 2445 |     get(tuple<_Elements...>& __t) noexcept
      |     ^~~
/usr/include/c++/14/bits/stl_pair.h:1265:5: note:                 'template<long unsigned int _Int, class _Tp1, class _Tp2> constexpr const typename std::tuple_element<_Int, std::pair<_Tp1, _Tp2> >::type&& std::get(const pair<_Tp1, _Tp2>&&)'
 1265 |     get(const pair<_Tp1, _Tp2>&& __in) noexcept
      |     ^~~
/usr/include/c++/14/bits/stl_pair.h:1260:5: note:                 'template<long unsigned int _Int, class _Tp1, class _Tp2> constexpr const typename std::tuple_element<_Int, std::pair<_Tp1, _Tp2> >::type& std::get(const pair<_Tp1, _Tp2>&)'
 1260 |     get(const pair<_Tp1, _Tp2>& __in) noexcept
      |     ^~~
/usr/include/c++/14/bits/stl_pair.h:1255:5: note:                 'template<long unsigned int _Int, class _Tp1, class _Tp2> constexpr typename std::tuple_element<_Int, std::pair<_Tp1, _Tp2> >::type&& std::get(pair<_Tp1, _Tp2>&&)'
 1255 |     get(pair<_Tp1, _Tp2>&& __in) noexcept
      |     ^~~
/usr/include/c++/14/bits/stl_pair.h:1250:5: note:                 'template<long unsigned int _Int, class _Tp1, class _Tp2> constexpr typename std::tuple_element<_Int, std::pair<_Tp1, _Tp2> >::type& std::get(pair<_Tp1, _Tp2>&)'
 1250 |     get(pair<_Tp1, _Tp2>& __in) noexcept
      |     ^~~
a.cc:13:12: note:                 'std::vector<long long int> get'
   13 | vector<ll> get; //各DFSで見た部分木に含まれる頂点のリスト
      |            ^~~
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 function 'void dfs3(ll, ll)':
a.cc:47:5: error: reference to 'get' is ambiguous
   47 |     get.emplace_back(now);
      |     ^~~
/usr/include/c++/14/variant:1180:5: note: candidates are: 'template<class _Tp, class ... _Types> constexpr const _Tp&& std::get(const variant<_Types ...>&&)'
 1180 |     get(const variant<_Types...>&& __v)
      |     ^~~
/usr/include/c++/14/variant:1170:5: note:                 'template<class _Tp, class ... _Types> constexpr const _Tp& std::get(const variant<_Types ...>&)'
 1170 |     get(const variant<_Types...>& __v)
      |     ^~~
/usr/include/c++/14/variant:1160:5: note:                 'template<class _Tp, class ... _Types> constexpr _Tp&& std::get(variant<_Types ...>&&)'
 1160 |     get(variant<_Types...>&& __v)
      |     ^~~
/usr/include/c++/14/variant:1150:5: note:                 'template<class _Tp, class ... _Types> constexpr _Tp& std::get(variant<_Types ...>&)'
 1150 |     get(variant<_Types...>& __v)