Language: C++

#include <iostream>
#include <string>
#include <vector>
#include <algorithm>
#include <sstream>
#include <queue>
#include <deque>
#include <bitset>
#include <iterator>
#include <list>
#include <stack>
#include <map>
#include <set>
#include <functional>
#include <numeric>
#include <utility>
#include <iomanip>
#include <limits>
#include <time.h>
#include <math.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <assert.h>

#define debug(x) cout << #x << " = " << x << endl
#define fori(i, ini, lim) for(int i = int(ini); i < int(lim); i++)
#define ford(i, ini, lim) for(int i = int(ini); i >= int(lim); i--)

using namespace std;

typedef long long ll;
typedef long double ld;
typedef pair<int, int> ii;

const int MAX = 2e5 + 5;
map<int, int> data[MAX];
vector<int> adj[MAX];
ll ans[MAX];
int color[MAX];
int n;

int dfs(int source, int father) {
    int size = 1;
    int c = color[source];
    for (auto &each : adj[source]) {
        if (each == father) continue;
        int each_size = dfs(each, source);
        size += each_size;
        int not_my_color = each_size - data[each][c];
        ans[c] += (1LL * not_my_color * (not_my_color + 1)) / 2;
        if (data[each].size() > data[source].size()) swap(data[source], data[each]);
        for (auto &d : data[each]) {
            data[source][d.first] += d.second;    
        }
    }
    int not_my_color = (size - 1) - data[source][c];
    data[source][c] += not_my_color + 1;
    return size;
}

void solve() {
    cin >> n;
    fori (i, 1, n + 1) cin >> color[i];
    fori (i, 0, n - 1) {
        int u, v;
        cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    dfs(1, -1);
    fori (i, 1, n + 1) {
        int c = i;
        if (c == color[1]) continue;
        int not_my_color = n - data[1][c];
        ans[c] += (1LL * not_my_color * (not_my_color + 1)) / 2;
    }
    ll total_paths = (1LL * n * (n + 1)) / 2;
    fori (i, 1, n + 1) cout << total_paths - ans[i] << "\n";
}

int main() {
    ios_base::sync_with_stdio(false);
    solve();
    return 0;
}



Compiler output:
a.cc: In function 'int dfs(int, int)':
a.cc:50:40: error: reference to 'data' is ambiguous
   50 |         int not_my_color = each_size - data[each][c];
      |                                        ^~~~
In file included from /usr/include/c++/14/string:53,
                 from /usr/include/c++/14/bits/locale_classes.h:40,
                 from /usr/include/c++/14/bits/ios_base.h:41,
                 from /usr/include/c++/14/ios:44,
                 from /usr/include/c++/14/ostream:40,
                 from /usr/include/c++/14/iostream:41,
                 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:37:15: note:                 'std::map<int, int> data [200005]'
   37 | map<int, int> data[MAX];
      |               ^~~~
a.cc:52:13: error: reference to 'data' is ambiguous
   52 |         if (data[each].size() > data[source].size()) swap(data[source], data[each]);
      |             ^~~~
/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:37:15: note:                 'std::map<int, int> data [200005]'
   37 | map<int, int> data[MAX];
      |               ^~~~
a.cc:52:33: error: reference to 'data' is ambiguous
   52 |         if (data[each].size() > data[source].size()) swap(data[source], data[each]);
      |                                 ^~~~
/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:37:15: note:                 'std::map<int, int> data [200005]'
   37 | map<int, int> data[MAX];
      |               ^~~~
a.cc:52:59: error: reference to 'data' is ambiguous
   52 |         if (data[each].size() > data[source].size()) swap(data[source], data[each]);
      |                                                           ^~~~
/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:37:15: note:                 'std::map<int, int> data [200005]'
   37 | map<int, int> data[MAX];
      |               ^~~~
a.cc:52:73: error: reference to 'data' is ambiguous
   52 |         if (data[each].size() > data[source].size()) swap(data[source], data[each]);
      |                                                                         ^~~~
/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:37:15: note:                 'std::map<int, int> data [200005]'
   37 | map<int, int> data[MAX];
      |               ^~~~
a.cc:53:24: error: reference to 'data' is ambiguous
   53 |         for (auto &d : data[each]) {
      |                        ^~~~
/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:37:15: note:                 'std::map<int, int> data [200005]'
   37 | map<int, int> data[MAX];
      |               ^~~~
a.cc:54:13: error: reference to 'data' is ambiguous
   54 |             data[source][d.first] += d.second;
      |             ^~~~
/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:37:15: note:                 'std::map<int, int> data [200005]'
   37 | map<int, int> data[MAX];
      |               ^~~~
a.cc:57:37: error: reference to 'data' is ambiguous
   57 |     int not_my_color = (size - 1) - data[source][c];
      |                                     ^~~~
/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:37:15: note:                 'std::map<int, int> data [2000