Language: C++

#include <bits/stdc++.h>
#define RI register int
typedef long long LL;
#define int LL

#define FILEIO(name) freopen(name".in", "r", stdin), freopen(name".out", "w", stdout);

using namespace std;

char buf[1000000], *p1 = buf, *p2 = buf;
inline char gc() {
  if (p1 == p2) p2 = (p1 = buf) + fread(buf, 1, 1000000, stdin);
  return p1 == p2 ? EOF : *(p1++);
}
template <class T> inline void read(T &n) {
  n = 0; RI ch = gc(), f;
  while ((ch < '0' || ch > '9') && ch != '-') ch = gc();
  f = (ch == '-' ? ch = gc(), -1 : 1);
  while (ch >= '0' && ch <= '9') n = n * 10 + (ch ^ 48), ch = gc();
  n *= f;
}

int const MAXN = 2e5 + 5;
struct Edges { int to, next; } e[MAXN * 2];
int head[MAXN], tot;
int Ans[MAXN], size[MAXN], Out[MAXN];
int a[MAXN];

inline void addedge(int from, int to) {
  e[++tot] = (Edges){to, head[from]};
  head[from] = tot;
  e[++tot] = (Edges){from, head[to]};
  head[to] = tot;
}

void Dfs(int now, int fa) {
  size[now] = 1;
  for (RI i = head[now]; i; i = e[i].next)
    if (e[i].to != fa) {
      int nowsplit = Out[a[now]];
      Dfs(e[i].to, now);
      size[now] += size[e[i].to];
      int addsize = Out[a[now]] - nowsplit;
      int S = size[e[i].to] - addsize;
      Out[a[now]] += S;
      Ans[a[now]] -= S * (S + 1) / 2;
    }
  ++Out[a[now]];
}

signed main() {
  
#ifdef LOCAL
  FILEIO("a");
#endif

  int n; read(n);
  for (RI i = 1; i <= n; ++i)
    read(a[i]), Ans[i] = n * (n + 1) / 2;
  for (RI i = 1, x, y; i < n; ++i)
    read(x), read(y), addedge(x, y);
  Dfs(1, 0);
  for (RI i = 1; i <= n; ++i) {
    Ans[i] -= (n - Out[i]) * (n - Out[i] + 1) / 2;
    printf("%lld\n", Ans[i]);
  }

  return 0;
}

// created by Daniel yuan


Compiler output:
a.cc: In function 'void read(T&)':
a.cc:16:13: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
   16 |   n = 0; RI ch = gc(), f;
      |             ^~
a.cc:16:24: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
   16 |   n = 0; RI ch = gc(), f;
      |                        ^
a.cc: In function 'void Dfs(LL, LL)':
a.cc:37:3: error: reference to 'size' is ambiguous
   37 |   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:26:16: note:                 'LL size [200005]'
   26 | int Ans[MAXN], size[MAXN], Out[MAXN];
      |                ^~~~
a.cc:38:11: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
   38 |   for (RI i = head[now]; i; i = e[i].next)
      |           ^
a.cc:42:7: error: reference to 'size' is ambiguous
   42 |       size[now] += size[e[i].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:26:16: note:                 'LL size [200005]'
   26 | int Ans[MAXN], size[MAXN], Out[MAXN];
      |                ^~~~
a.cc:42:20: error: reference to 'size' is ambiguous
   42 |       size[now] += size[e[i].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:26:16: note:                 'LL size [200005]'
   26 | int Ans[MAXN], size[MAXN], Out[MAXN];
      |                ^~~~
a.cc:44:15: error: reference to 'size' is ambiguous
   44 |       int S = size[e[i].to] - addsize;
      |               ^~~~
/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:26:16: note:                 'LL size [200005]'
   26 | int Ans[MAXN], size[MAXN], Out[MAXN];
      |                ^~~~
a.cc: In function 'int main()':
a.cc:58:11: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
   58 |   for (RI i = 1; i <= n; ++i)
      |           ^
a.cc:60:11: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
   60 |   for (RI i = 1, x, y; i < n; ++i)
      |           ^
a.cc:60:18: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
   60 |   for (RI i = 1, x, y; i < n; ++i)
      |                  ^
a.cc:60:21: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
   60 |   for (RI i = 1, x, y; i < n; ++i)
      |                     ^
a.cc:63:11: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
   63 |   for (RI i = 1; i <= n; ++i) {
      |           ^
a.cc: In instantiation of 'void read(T&) [with T = long long int]':
a.cc:57:14:   required from here
   57 |   int n; read(n);
      |          ~~~~^~~
a.cc:16:13: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
   16 |   n = 0; RI ch = gc(), f;
      |             ^~
a.cc:16:24: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
   16 |   n = 0; RI ch = gc(), f;
      |                        ^