Language: C++

#include <bits/stdc++.h>
#define R register
#define mp make_pair
#define ll long long
#define pii pair<int, int>
using namespace std;
const int mod = 998244353, N = 210000;

int n, c[N], dfn[N], size[N], top[N], son[N], fa[N], dep[N], cnt, f[N], stck[N], tp;
ll g[N], ans;
vector<int> to[N], pos[N];
vector<pii> ito[N];

inline int addMod(int a, int b) {
	return (a += b) >= mod ? a - mod : a;
}

inline ll quickpow(ll base, ll pw) {
	ll ret = 1;
	while (pw) {
		if (pw & 1) ret = ret * base % mod;
		base = base * base % mod, pw >>= 1;
	}
	return ret;
}

template <class T>
inline void read(T &x) {
	x = 0;
	char ch = getchar(), w = 0;
	while (!isdigit(ch)) w = (ch == '-'), ch = getchar();
	while (isdigit(ch)) x = (x << 1) + (x << 3) + (ch ^ 48), ch = getchar();
	x = w ? -x : x;
	return;
}

ll comb(ll x) {
	return x * (x + 1) / 2;
}

inline void addEdg(int x, int y) {
	to[x].push_back(y), to[y].push_back(x);
	return;
}

void dfs1(int now) {
	size[now] = 1, dep[now] = dep[fa[now]] + 1;
	for (auto &v : to[now]) {
		if (v == fa[now]) continue;
		fa[v] = now, dfs1(v), size[now] += size[v];
		g[now] += comb(size[v]);
		if (size[son[now]] < size[v]) son[now] = v;
	}
	return;
}

void dfs2(int now) {
	dfn[now] = ++cnt;
	top[now] = now == son[fa[now]] ? top[fa[now]] : now;
	for (auto &v : to[now]) {
		if (v == fa[now]) continue;
		dfs2(v);
	}
	return;
}

inline int getLca(int x, int y) {
	while (top[x] != top[y])
		dep[top[x]] >= dep[top[y]] ? x = fa[top[x]] : y = fa[top[y]];
	return dep[x] >= dep[y] ? y : x;
}

inline int getSon(int x, int anc) {
	if (top[x] == top[anc]) return son[anc];
	x = top[x];
	while (top[fa[x]] != top[anc]) x = top[fa[x]];
	return fa[x] == anc ? x : son[anc];
}

inline bool cmp(int x, int y) {
	return dfn[x] < dfn[y];
}

inline void insrt(int x) {
	stck[++tp] = x, ito[x].clear(), f[x] = 0;
	return;
}

void addIEdg(int x, int y) {
	ito[x].emplace_back(y, getSon(y, x));
	return;
}

void dfs3(int now, int col) {
	if (c[now] == col) {
		ans += g[now];
		for (auto &v : ito[now]) {
			dfs3(v.first, col);
			ans -= comb(size[v.second]);
			ans += comb(size[v.second] - size[v.first] + f[v.first]);
		}
		return;
	}
	f[now] = size[now];
	for (auto &v : ito[now]) {
		dfs3(v.first, col);
		f[now] -= size[v.first] - f[v.first];
	}
	return;
}

int main() {
	int x, y;
	read(n);
	for (R int i = 1; i <= n; ++i)
		read(c[i]), pos[c[i]].push_back(i);
	for (R int i = 1; i < n; ++i)
		read(x), read(y), addEdg(x, y);
	dfs1(1), dfs2(1);
	for (R int i = 1; i <= n; ++i) {
		if (c[1] != i) pos[i].push_back(1);
		sort(pos[i].begin(), pos[i].end(), cmp);
		tp = 0, insrt(1);
		for (auto &v : pos[i]) {
			if (v == 1) continue;
			int lca = getLca(stck[tp], v);
			while (tp > 1 && lca != stck[tp]) {
				int tmp = stck[tp--];
				if (dep[lca] > dep[stck[tp]]) insrt(lca);
				addIEdg(stck[tp], tmp), lca = getLca(stck[tp], v);
			}
			insrt(v);
		}
		while (tp > 1) addIEdg(stck[tp - 1], stck[tp]), --tp;
		ans = 0, dfs3(1, i);
		printf("%lld\n", comb(n) - ans - comb(f[1]));
	}
	return 0;
}

Compiler output:
a.cc: In function 'void dfs1(int)':
a.cc:47:9: error: reference to 'size' is ambiguous
   47 |         size[now] = 1, dep[now] = dep[fa[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:9:22: note:                 'int size [210000]'
    9 | int n, c[N], dfn[N], size[N], top[N], son[N], fa[N], dep[N], cnt, f[N], stck[N], tp;
      |                      ^~~~
a.cc:50:39: error: reference to 'size' is ambiguous
   50 |                 fa[v] = now, dfs1(v), size[now] += size[v];
      |                                       ^~~~
/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:9:22: note:                 'int size [210000]'
    9 | int n, c[N], dfn[N], size[N], top[N], son[N], fa[N], dep[N], cnt, f[N], stck[N], tp;
      |                      ^~~~
a.cc:50:52: error: reference to 'size' is ambiguous
   50 |                 fa[v] = now, dfs1(v), size[now] += size[v];
      |                                                    ^~~~
/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:9:22: note:                 'int size [210000]'
    9 | int n, c[N], dfn[N], size[N], top[N], son[N], fa[N], dep[N], cnt, f[N], stck[N], tp;
      |                      ^~~~
a.cc:51:32: error: reference to 'size' is ambiguous
   51 |                 g[now] += comb(size[v]);
      |                                ^~~~
/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:9:22: note:                 'int size [210000]'
    9 | int n, c[N], dfn[N], size[N], top[N], son[N], fa[N], dep[N], cnt, f[N], stck[N], tp;
      |                      ^~~~
a.cc:52:21: error: reference to 'size' is ambiguous
   52 |                 if (size[son[now]] < size[v]) son[now] = v;
      |                     ^~~~
/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:9:22: note:                 'int size [210000]'
    9 | int n, c[N], dfn[N], size[N], top[N], son[N], fa[N], dep[N], cnt, f[N], stck[N], tp;
      |                      ^~~~
a.cc:52:38: error: reference to 'size' is ambiguous
   52 |                 if (size[son[now]] < size[v]) son[now] = v;
      |                                      ^~~~
/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:9:22: note:                 'int size [210000]'
    9 | int n, c[N], dfn[N], size[N], top[N], son[N], fa[N], dep[N], cnt, f[N], stck[N], tp;
      |                      ^~~~
a.cc: In function 'void dfs3(int, int)':
a.cc:99:37: error: reference to 'size' is ambiguous
   99 |                         ans -= comb(size[v.second]);
      |                                     ^~~~
/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:9:22: note:                 'int size [210000]'
    9 | int n, c[N], dfn[N], size[N], top[N], son[N], fa[N], dep[N], cnt, f[N], stck[N], tp;
      |                      ^~~~
a.cc:100:37: error: reference to 'size' is ambiguous
  100 |                         ans += comb(size[v.second] - size[v.first] + f[v.first]);
      |                                     ^~~~
/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:9:22: note:                 'int size [210000]'
    9 | int n, c[N], dfn[N], size[N], top[N], son[N], fa[N], dep[N], cnt, f[N], stck[N], tp;
      |                      ^~~~
a.cc:100:54: error: reference to 'size' is ambiguous
  100 |                         ans += comb(size[v.second] - size[v.first] + f[v.first]);
      |                                                      ^~~~
/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:9:22: note:                 'int size [210000]'
    9 | int n, c[N], dfn[N], size[N], top[N], son[N], fa[N], dep[N], cnt, f[N], stck[N], tp;
      |                      ^~~~
a.cc:104:18: error: reference to 'size' is ambiguous
  104 |         f[now] = size[now];
      |                  ^~~~
/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:9:22: note:                 'int size [210000]'
    9 | int n, c[N], dfn[N], size[N], top[N], son[N], fa[N], dep[N], cnt, f[N], stck[N], tp;
      |                      ^~~~
a.cc:107:27: error: reference to 'size' is ambiguous
  107 |                 f[now] -= size[v.first] - f[v.first];
      |                           ^~~~
/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:9:22: note:                 'int size [210000]'
    9 | int n, c[N], dfn[N], size[N], top[N], son[N], fa[N], dep[N], cnt, f[N], stck[N], tp;
      |                      ^~~~
a.cc: In function 'int main()':
a.cc:115:20: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
  115 |         for (R int i = 1; i <= n; ++i)
      |                    ^
a.cc:117:20: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
  117 |         for (R int i = 1; i < n; ++i)
      |                    ^
a.cc:120:20: warning: ISO C++17 does not allow 'register' storage class specifier [-Wregister]
  120 |         for (R int i = 1; i <= n; ++i) {
      |                    ^