Language: C++

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<ll,ll> pr;
typedef vector<ll> vc;
#define pb push_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++)
struct FastIO {FastIO() { cin.tie(0); ios::sync_with_stdio(0); }}fastiofastio;
#ifdef DEBUG_ 
#include "../program_contest_library/memo/dump.hpp"
#else
#define dump(...)
#endif
ll n;
bool used[200005]={0};
vector<ll> sz(200005), dead(200005);
vector<ll> size(200005);
vector<ll> g[200005];
ll dp[200005]={0},cs[200005];
ll ans[200005]={0},cp;
vc ns;
ll sizem;

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

void dfs1(ll now, ll pr){
    ns.push_back(now);
    used[cs[now]]=true;
    for(auto to: g[now]) if(to!=pr&&!dead[to]) {
        if(!used[cs[to]]) dp[cs[to]]+=size[to];
        dfs1(to,now);
    }
    used[cs[now]]=false;
}

int main() {
    // your code goes here
    ll aa,bb,gc;
    cin >> n;
    rep(i,1,n){
        cin >> cs[i-1];
        cs[i-1]--;
    }
    rep(i,1,n-1){
        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);
    };
    vector<ll> depth(n);
    sizem = n;
    auto centroid_decomposition = [&](auto &&self, ll root, ll d) -> void {
        ll c = find_centroid(root);
        dead[c] = true;
        depth[c] = d;
        cp = cs[c];
        //cout << "center : " << c << " " << sizem << " " << cp << endl;
        vc dm[200005];
        ll memo[200005]={0};
        rep(ii,0,n-1){
            dp[ii]=0;
            dm[ii].push_back(0);
            used[ii]=false;
        }
        for(auto to: g[c]) if(!dead[to]) {
            dfs2(to,c);
            ns.clear();
            dp[cs[to]]+=size[to];
            used[cs[to]]=true;
            dfs1(to,c);
            //if(c==5&&to==4) cout << size[4] << endl;
            used[cs[to]]=false;
            ll usc[200005]={0};
            for(auto in: ns){
                if(!used[cs[in]]) dm[cs[in]].push_back(dp[cs[in]]);
                if(!used[cs[in]]) memo[cs[in]]+=( dm[cs[in]][dm[cs[in]].size()-1] - dm[cs[in]][dm[cs[in]].size()-2] )*(sizem-size[to]);
                used[cs[in]]=true;
            }
            dm[cp].pb(dp[cp]);
            for(auto in: ns){
                used[cs[in]]=false;
            }
        }
        rep(ii,0,n-1){
            if(ii==cp){
            	memo[ii]=0;
                for(auto to: g[c]){
                    if(dead[to]) continue;
                	ans[cp]+=size[to];
                    memo[cp]+=size[to]*(sizem-1-size[to]);
                }
                ans[cp]+=(memo[cp]/2);
                //cout << "ans3 : " << ii << " " << ans[ii] << endl;
                continue;
            }
            ans[ii] += memo[ii]; 
            //cout << "ans1 : " << ii << " " << ans[ii] << endl;
            memo[ii] = 0;
        	//cout << "dp : " << ii << " " << dp[ii] << endl;
            //cout << ii << " " << size[ii] << endl;
            memo[ii]-=dp[ii]*dp[ii];
            reps(jj,dm[ii]){
                if(jj==0) continue;
                memo[ii]+=(dm[ii][jj]-dm[ii][jj-1])*(dm[ii][jj]-dm[ii][jj-1]);
            }
            ans[ii]+=(memo[ii]/2);
            //cout << "ans2 : " << ii << " " << ans[ii] << endl;
        }
        used[cp] = false;
        for(auto to: g[c]) if(!dead[to]) {
            dump(root, c, to, d);
            sizem = size[to];
            self(self, to, d+1);
        }
        dead[c] = false;
    };
    centroid_decomposition(centroid_decomposition, 0, 0);
    rep(i,0,n-1) used[i]=false;
    rep(i,0,n-1) ans[cs[i]]++;
    rep(i,0,n-1) cout << ans[i] << endl;
    return 0;
}

Compiler output:
a.cc: In function 'void dfs2(ll, ll)':
a.cc:27:5: error: reference to 'size' is ambiguous
   27 |     size[now]=1;
      |     ^~~~
In file included from /usr/include/c++/14/string:53,
                 from /usr/include/c++/14/bitset:52,
                 from /usr/include/x86_64-linux-gnu/c++/14/bits/stdc++.h:52,
                 from a.cc:1:
/usr/include/c++/14/bits/range_access.h:272:5: note: candidates are: 'template<class _Tp, long unsigned int _Nm> constexpr std::size_t std::size(const _Tp (&)[_Nm])'
  272 |     size(const _Tp (&)[_Nm]) noexcept
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:262:5: note:                 'template<class _Container> constexpr decltype (__cont.size()) std::size(const _Container&)'
  262 |     size(const _Container& __cont) noexcept(noexcept(__cont.size()))
      |     ^~~~
a.cc:19:12: note:                 'std::vector<long long int> size'
   19 | vector<ll> size(200005);
      |            ^~~~
a.cc:30:9: error: reference to 'size' is ambiguous
   30 |         size[now] += size[to];
      |         ^~~~
/usr/include/c++/14/bits/range_access.h:272:5: note: candidates are: 'template<class _Tp, long unsigned int _Nm> constexpr std::size_t std::size(const _Tp (&)[_Nm])'
  272 |     size(const _Tp (&)[_Nm]) noexcept
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:262:5: note:                 'template<class _Container> constexpr decltype (__cont.size()) std::size(const _Container&)'
  262 |     size(const _Container& __cont) noexcept(noexcept(__cont.size()))
      |     ^~~~
a.cc:19:12: note:                 'std::vector<long long int> size'
   19 | vector<ll> size(200005);
      |            ^~~~
a.cc:30:22: error: reference to 'size' is ambiguous
   30 |         size[now] += size[to];
      |                      ^~~~
/usr/include/c++/14/bits/range_access.h:272:5: note: candidates are: 'template<class _Tp, long unsigned int _Nm> constexpr std::size_t std::size(const _Tp (&)[_Nm])'
  272 |     size(const _Tp (&)[_Nm]) noexcept
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:262:5: note:                 'template<class _Container> constexpr decltype (__cont.size()) std::size(const _Container&)'
  262 |     size(const _Container& __cont) noexcept(noexcept(__cont.size()))
      |     ^~~~
a.cc:19:12: note:                 'std::vector<long long int> size'
   19 | vector<ll> size(200005);
      |            ^~~~
a.cc: In function 'void dfs1(ll, ll)':
a.cc:38:39: error: reference to 'size' is ambiguous
   38 |         if(!used[cs[to]]) dp[cs[to]]+=size[to];
      |                                       ^~~~
/usr/include/c++/14/bits/range_access.h:272:5: note: candidates are: 'template<class _Tp, long unsigned int _Nm> constexpr std::size_t std::size(const _Tp (&)[_Nm])'
  272 |     size(const _Tp (&)[_Nm]) noexcept
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:262:5: note:                 'template<class _Container> constexpr decltype (__cont.size()) std::size(const _Container&)'
  262 |     size(const _Container& __cont) noexcept(noexcept(__cont.size()))
      |     ^~~~
a.cc:19:12: note:                 'std::vector<long long int> size'
   19 | vector<ll> size(200005);
      |            ^~~~
a.cc: In lambda function:
a.cc:93:25: error: reference to 'size' is ambiguous
   93 |             dp[cs[to]]+=size[to];
      |                         ^~~~
/usr/include/c++/14/bits/range_access.h:272:5: note: candidates are: 'template<class _Tp, long unsigned int _Nm> constexpr std::size_t std::size(const _Tp (&)[_Nm])'
  272 |     size(const _Tp (&)[_Nm]) noexcept
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:262:5: note:                 'template<class _Container> constexpr decltype (__cont.size()) std::size(const _Container&)'
  262 |     size(const _Container& __cont) noexcept(noexcept(__cont.size()))
      |     ^~~~
a.cc:19:12: note:                 'std::vector<long long int> size'
   19 | vector<ll> size(200005);
      |            ^~~~
a.cc:101:126: error: reference to 'size' is ambiguous
  101 |                 if(!used[cs[in]]) memo[cs[in]]+=( dm[cs[in]][dm[cs[in]].size()-1] - dm[cs[in]][dm[cs[in]].size()-2] )*(sizem-size[to]);
      |                                                                                                                              ^~~~
/usr/include/c++/14/bits/range_access.h:272:5: note: candidates are: 'template<class _Tp, long unsigned int _Nm> constexpr std::size_t std::size(const _Tp (&)[_Nm])'
  272 |     size(const _Tp (&)[_Nm]) noexcept
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:262:5: note:                 'template<class _Container> constexpr decltype (__cont.size()) std::size(const _Container&)'
  262 |     size(const _Container& __cont) noexcept(noexcept(__cont.size()))
      |     ^~~~
a.cc:19:12: note:                 'std::vector<long long int> size'
   19 | vector<ll> size(200005);
      |            ^~~~
a.cc:114:34: error: reference to 'size' is ambiguous
  114 |                         ans[cp]+=size[to];
      |                                  ^~~~
/usr/include/c++/14/bits/range_access.h:272:5: note: candidates are: 'template<class _Tp, long unsigned int _Nm> constexpr std::size_t std::size(const _Tp (&)[_Nm])'
  272 |     size(const _Tp (&)[_Nm]) noexcept
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:262:5: note:                 'template<class _Container> constexpr decltype (__cont.size()) std::size(const _Container&)'
  262 |     size(const _Container& __cont) noexcept(noexcept(__cont.size()))
      |     ^~~~
a.cc:19:12: note:                 'std::vector<long long int> size'
   19 | vector<ll> size(200005);
      |            ^~~~
a.cc:115:31: error: reference to 'size' is ambiguous
  115 |                     memo[cp]+=size[to]*(sizem-1-size[to]);
      |                               ^~~~
/usr/include/c++/14/bits/range_access.h:272:5: note: candidates are: 'template<class _Tp, long unsigned int _Nm> constexpr std::size_t std::size(const _Tp (&)[_Nm])'
  272 |     size(const _Tp (&)[_Nm]) noexcept
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:262:5: note:                 'template<class _Container> constexpr decltype (__cont.size()) std::size(const _Container&)'
  262 |     size(const _Container& __cont) noexcept(noexcept(__cont.size()))
      |     ^~~~
a.cc:19:12: note:                 'std::vector<long long int> size'
   19 | vector<ll> size(200005);
      |            ^~~~
a.cc:115:49: error: reference to 'size' is ambiguous
  115 |                     memo[cp]+=size[to]*(sizem-1-size[to]);
      |                                                 ^~~~
/usr/include/c++/14/bits/range_access.h:272:5: note: candidates are: 'template<class _Tp, long unsigned int _Nm> constexpr std::size_t std::size(const _Tp (&)[_Nm])'
  272 |     size(const _Tp (&)[_Nm]) noexcept
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:262:5: note:                 'template<class _Container> constexpr decltype (__cont.size()) std::size(const _Container&)'
  262 |     size(const _Container& __cont) noexcept(noexcept(__cont.size()))
      |     ^~~~
a.cc:19:12: note:                 'std::vector<long long int> size'
   19 | vector<ll> size(200005);
      |            ^~~~
a.cc:137:21: error: reference to 'size' is ambiguous
  137 |             sizem = size[to];
      |                     ^~~~
/usr/include/c++/14/bits/range_access.h:272:5: note: candidates are: 'template<class _Tp, long unsigned int _Nm> constexpr std::size_t std::size(const _Tp (&)[_Nm])'
  272 |     size(const _Tp (&)[_Nm]) noexcept
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:262:5: note:                 'template<class _Container> constexpr decltype (__cont.size()) std::size(const _Container&)'
  262 |     size(const _Container& __cont) noexcept(noexcept(__cont.size()))
      |     ^~~~
a.cc:19:12: note:                 'std::vector<long long int> size'
   19 | vector<ll> size(200005);
      |            ^~~~