Language: C++

#include <bits/stdc++.h>
using namespace std;

vector<vector<int>> graph;
vector<int> dive, ascend;// subtree_size info

int k = 0;
void DFS(int v, int parent_v = -1){
    dive[v] = k;
    k += 1;

    for(int child_v : graph[v]){
        if(child_v == parent_v) continue;
        DFS(child_v, v);
    }
    ascend[v] = k;
}

bool compare(int a, int b){
    return dive[a] > dive[b];
}

vector<int> data;
struct BIT{
    void init(int n){
        data.assign(n+1, 0);// 1-indexed [1,n]
    }

    void add(int i, int x){
        while(i < data.size()){
            data[i] += x;
            i += (i & -i);
        }
    }

    int sum(int i){
        int s = 0;
        while(i > 0){
            s += data[i];
            i -= (i & -i);
        }
        return s;
    }
};

int main(){
    long long int n;
    cin >> n;

    graph.resize(n);
    dive.resize(n);
    ascend.resize(n);

    vector<vector<int>> color(n);
    for(int i=0; i<n; i++){
        int c;
        cin >> c;
        c -= 1;
        color[c].push_back(i);
    }

    for(int i=0; i<n-1; i++){
        int a, b;
        cin >> a >> b;
        a -= 1;
        b -= 1;
        graph[a].push_back(b);
        graph[b].push_back(a);
    }
    DFS(0);
    //for(int i=0; i<n; i++) cout << dive[i] << " ";
    //cout << endl;
    //for(int i=0; i<n; i++) cout << ascend[i] << " ";
    //cout << endl;

    BIT bit; 
    bit.init(n);
    for(int i=0; i<n; i++){
        bit.add(i+1, 1);
    }

    for(int i=0; i<n; i++){
        long long int ans = n * (n-1) / 2 + n;

        sort(color[i].begin(), color[i].end(), compare);// deeper nodes first
        vector<pair<int, int>> rec;

        for(int v : color[i]){
            int count = 1;// subtraction in BIT

            for(int child_v : graph[v]){
                if(dive[child_v] < dive[v]) continue;

                long long int subtree_size = bit.sum(ascend[child_v]) - bit.sum(dive[child_v]);
                ans   -= subtree_size * (subtree_size - 1) / 2 + subtree_size;
                count += subtree_size;
            }

            bit.add(dive[v]+1, -count);
            rec.push_back({dive[v]+1, count});
        }
        long long int subtree_size = bit.sum(n);
        ans -= subtree_size * (subtree_size - 1) / 2 + subtree_size;
        cout << ans << endl;

        for(auto r : rec) bit.add(r.first, r.second);// avoid TLE due to "for(int i=0; i<n; i++) bit.add(i,1)" everytime
    }
    return 0;
}

Compiler output:
a.cc: In member function 'void BIT::init(int)':
a.cc:26:9: error: reference to 'data' is ambiguous
   26 |         data.assign(n+1, 0);// 1-indexed [1,n]
      |         ^~~~
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:23:13: note:                 'std::vector<int> data'
   23 | vector<int> data;
      |             ^~~~
a.cc: In member function 'void BIT::add(int, int)':
a.cc:30:19: error: reference to 'data' is ambiguous
   30 |         while(i < data.size()){
      |                   ^~~~
/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:23:13: note:                 'std::vector<int> data'
   23 | vector<int> data;
      |             ^~~~
a.cc:31:13: error: reference to 'data' is ambiguous
   31 |             data[i] += x;
      |             ^~~~
/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:23:13: note:                 'std::vector<int> data'
   23 | vector<int> data;
      |             ^~~~
a.cc: In member function 'int BIT::sum(int)':
a.cc:39:18: error: reference to 'data' is ambiguous
   39 |             s += 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:23:13: note:                 'std::vector<int> data'
   23 | vector<int> data;
      |             ^~~~