Language: C++

#include <bits/stdc++.h>
using namespace std;
#define FOR(i, a, b) for(int i=(a);i<(b);++i)
#define rep(i, n) FOR(i, 0, n)
#define whole(x) (x).begin(),(x).end()
#define UNIQUE(v) v.erase(unique(v.begin(), v.end()), v.end())
using ll = long long;
using P = pair<int, int>;
const int mod = 1000000007;

vector<int> to[200005];
int c[200005];
int buka[200005];
vector<int> data[200005];
int numcol[200005];

void dfs(int x, int p) {
    if (to[x].size()==1 && to[x][0]==p) return;
    for (int e: to[x]) {
        if (e==p) continue;
        dfs(e, x);
        buka[x] += (buka[e])+1;
    }
}

void dfs2(int x, int p) {
    int color = c[x];
    if (data[color].size()) {
        data[color][data[color].size()-1] -= (buka[x]+1);
    }
    int id = data[color].size()-1;
    for (int e: to[x]) {
        if (e==p) continue;
        data[color].push_back(buka[e]+1);
        //int v = data[color][id];
        //data[color].erase(data[color].begin()+id);
        //data[color].push_back(v);
        //id = data[color].size()-1;
        dfs2(e, x);
    }
}

int main(){
    int n;
    cin >> n;
    
    
    rep(i, n) {
        cin >> c[i];
        c[i]--;
        numcol[c[i]]++;
    }
    rep(i, n-1) {
        int a, b;
        cin >> a >> b;
        a--; b--;
        to[a].emplace_back(b);
        to[b].emplace_back(a);
    }

    dfs(0, -1);
    dfs2(0, -1);

    rep(i, n) {
        /*
        rep(j, data[i].size()) {
            cout << data[i][j] << " ";
        }
        */
        //cout << endl;
        if (data[i].size()==0 && numcol[i]==0) {
            //ll ans2 = numcol[i]*(numcol[i]+1)/2;
            cout << "0" << endl;
            continue;
        }
        ll sum = 0;
        ll ans = 0;
        for (int e: data[i]) {
            ll x = e;
            sum += x;
            ans += x*(x+1)/2;
        }
        ll rem = n - sum - numcol[i];
        ans += rem*(rem+1)/2;
        ll res = n*(n+1)/2 - ans;
        cout << res << endl;
    }

    
    
    
    
    return 0;
}


Compiler output:
a.cc: In function 'void dfs2(int, int)':
a.cc:28:9: error: reference to 'data' is ambiguous
   28 |     if (data[color].size()) {
      |         ^~~~
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:344:5: note: candidates are: 'template<class _Tp> constexpr const _Tp* std::data(initializer_list<_Tp>)'
  344 |     data(initializer_list<_Tp> __il) noexcept
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:334:5: note:                 'template<class _Tp, long unsigned int _Nm> constexpr _Tp* std::data(_Tp (&)[_Nm])'
  334 |     data(_Tp (&__array)[_Nm]) noexcept
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:323:5: note:                 'template<class _Container> constexpr decltype (__cont.data()) std::data(const _Container&)'
  323 |     data(const _Container& __cont) noexcept(noexcept(__cont.data()))
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:312:5: note:                 'template<class _Container> constexpr decltype (__cont.data()) std::data(_Container&)'
  312 |     data(_Container& __cont) noexcept(noexcept(__cont.data()))
      |     ^~~~
a.cc:14:13: note:                 'std::vector<int> data [200005]'
   14 | vector<int> data[200005];
      |             ^~~~
a.cc:29:9: error: reference to 'data' is ambiguous
   29 |         data[color][data[color].size()-1] -= (buka[x]+1);
      |         ^~~~
/usr/include/c++/14/bits/range_access.h:344:5: note: candidates are: 'template<class _Tp> constexpr const _Tp* std::data(initializer_list<_Tp>)'
  344 |     data(initializer_list<_Tp> __il) noexcept
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:334:5: note:                 'template<class _Tp, long unsigned int _Nm> constexpr _Tp* std::data(_Tp (&)[_Nm])'
  334 |     data(_Tp (&__array)[_Nm]) noexcept
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:323:5: note:                 'template<class _Container> constexpr decltype (__cont.data()) std::data(const _Container&)'
  323 |     data(const _Container& __cont) noexcept(noexcept(__cont.data()))
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:312:5: note:                 'template<class _Container> constexpr decltype (__cont.data()) std::data(_Container&)'
  312 |     data(_Container& __cont) noexcept(noexcept(__cont.data()))
      |     ^~~~
a.cc:14:13: note:                 'std::vector<int> data [200005]'
   14 | vector<int> data[200005];
      |             ^~~~
a.cc:29:21: error: reference to 'data' is ambiguous
   29 |         data[color][data[color].size()-1] -= (buka[x]+1);
      |                     ^~~~
/usr/include/c++/14/bits/range_access.h:344:5: note: candidates are: 'template<class _Tp> constexpr const _Tp* std::data(initializer_list<_Tp>)'
  344 |     data(initializer_list<_Tp> __il) noexcept
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:334:5: note:                 'template<class _Tp, long unsigned int _Nm> constexpr _Tp* std::data(_Tp (&)[_Nm])'
  334 |     data(_Tp (&__array)[_Nm]) noexcept
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:323:5: note:                 'template<class _Container> constexpr decltype (__cont.data()) std::data(const _Container&)'
  323 |     data(const _Container& __cont) noexcept(noexcept(__cont.data()))
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:312:5: note:                 'template<class _Container> constexpr decltype (__cont.data()) std::data(_Container&)'
  312 |     data(_Container& __cont) noexcept(noexcept(__cont.data()))
      |     ^~~~
a.cc:14:13: note:                 'std::vector<int> data [200005]'
   14 | vector<int> data[200005];
      |             ^~~~
a.cc:31:14: error: reference to 'data' is ambiguous
   31 |     int id = data[color].size()-1;
      |              ^~~~
/usr/include/c++/14/bits/range_access.h:344:5: note: candidates are: 'template<class _Tp> constexpr const _Tp* std::data(initializer_list<_Tp>)'
  344 |     data(initializer_list<_Tp> __il) noexcept
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:334:5: note:                 'template<class _Tp, long unsigned int _Nm> constexpr _Tp* std::data(_Tp (&)[_Nm])'
  334 |     data(_Tp (&__array)[_Nm]) noexcept
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:323:5: note:                 'template<class _Container> constexpr decltype (__cont.data()) std::data(const _Container&)'
  323 |     data(const _Container& __cont) noexcept(noexcept(__cont.data()))
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:312:5: note:                 'template<class _Container> constexpr decltype (__cont.data()) std::data(_Container&)'
  312 |     data(_Container& __cont) noexcept(noexcept(__cont.data()))
      |     ^~~~
a.cc:14:13: note:                 'std::vector<int> data [200005]'
   14 | vector<int> data[200005];
      |             ^~~~
a.cc:34:9: error: reference to 'data' is ambiguous
   34 |         data[color].push_back(buka[e]+1);
      |         ^~~~
/usr/include/c++/14/bits/range_access.h:344:5: note: candidates are: 'template<class _Tp> constexpr const _Tp* std::data(initializer_list<_Tp>)'
  344 |     data(initializer_list<_Tp> __il) noexcept
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:334:5: note:                 'template<class _Tp, long unsigned int _Nm> constexpr _Tp* std::data(_Tp (&)[_Nm])'
  334 |     data(_Tp (&__array)[_Nm]) noexcept
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:323:5: note:                 'template<class _Container> constexpr decltype (__cont.data()) std::data(const _Container&)'
  323 |     data(const _Container& __cont) noexcept(noexcept(__cont.data()))
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:312:5: note:                 'template<class _Container> constexpr decltype (__cont.data()) std::data(_Container&)'
  312 |     data(_Container& __cont) noexcept(noexcept(__cont.data()))
      |     ^~~~
a.cc:14:13: note:                 'std::vector<int> data [200005]'
   14 | vector<int> data[200005];
      |             ^~~~
a.cc: In function 'int main()':
a.cc:71:13: error: reference to 'data' is ambiguous
   71 |         if (data[i].size()==0 && numcol[i]==0) {
      |             ^~~~
/usr/include/c++/14/bits/range_access.h:344:5: note: candidates are: 'template<class _Tp> constexpr const _Tp* std::data(initializer_list<_Tp>)'
  344 |     data(initializer_list<_Tp> __il) noexcept
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:334:5: note:                 'template<class _Tp, long unsigned int _Nm> constexpr _Tp* std::data(_Tp (&)[_Nm])'
  334 |     data(_Tp (&__array)[_Nm]) noexcept
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:323:5: note:                 'template<class _Container> constexpr decltype (__cont.data()) std::data(const _Container&)'
  323 |     data(const _Container& __cont) noexcept(noexcept(__cont.data()))
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:312:5: note:                 'template<class _Container> constexpr decltype (__cont.data()) std::data(_Container&)'
  312 |     data(_Container& __cont) noexcept(noexcept(__cont.data()))
      |     ^~~~
a.cc:14:13: note:                 'std::vector<int> data [200005]'
   14 | vector<int> data[200005];
      |             ^~~~
a.cc:78:21: error: reference to 'data' is ambiguous
   78 |         for (int e: data[i]) {
      |                     ^~~~
/usr/include/c++/14/bits/range_access.h:344:5: note: candidates are: 'template<class _Tp> constexpr const _Tp* std::data(initializer_list<_Tp>)'
  344 |     data(initializer_list<_Tp> __il) noexcept
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:334:5: note:                 'template<class _Tp, long unsigned int _Nm> constexpr _Tp* std::data(_Tp (&)[_Nm])'
  334 |     data(_Tp (&__array)[_Nm]) noexcept
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:323:5: note:                 'template<class _Container> constexpr decltype (__cont.data()) std::data(const _Container&)'
  323 |     data(const _Container& __cont) noexcept(noexcept(__cont.data()))
      |     ^~~~
/usr/include/c++/14/bits/range_access.h:312:5: note:                 'template<class _Container> constexpr decltype (__cont.data()) std::data(_Container&)'
  312 |     data(_Container& __cont) noexcept(noexcept(__cont.data()))
      |     ^~~~
a.cc:14:13: note:                 'std::vector<int> data [200005]'
   14 | vector<int> data[200005];
      |             ^~~~