Language: C++

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef unsigned long long ull;
typedef pair<int,int> pii;
typedef vector<int> vi; 
#define pb push_back
#define mp make_pair
#define fi first
#define se second
#define rep(i,a,b) for(int i=(a);i<=(b);i++)
#define per(i,a,b) for(int i=(a);i>=(b);i--)
const int mod = 1e9+7;
const int maxn =2e5+10;
const int inf = 0x3f3f3f3f;
vi g[maxn];
int id[maxn],cnt,siz[maxn],c[maxn],fa[maxn];
vi col[maxn];
void dfs(int now,int f) {
	siz[now]=1;
	fa[now]=f;
	id[now]=++cnt;
	for(int x:g[now]) {
		if(x==f) continue;
		dfs(x,now);
		siz[now]+=siz[x];
	}
}
ll data[maxn];
int lowbit(int x) {
	return x&(-x);
}
void add(int pos,int x) {
	for(int i=pos;i<maxn;i+=lowbit(i)) data[i]+=x;
}
ll query(int pos) {
	ll ans=1;
	for(int i=pos;i;i-=lowbit(i)) ans+=data[i];
	return ans;
}
ll n;
ll cal(int now) {
	return siz[now]+(query(id[now]+siz[now]-1)-query(id[now]-1));
}
ll b[maxn];
void solve(int cc) {
	if(col[cc].size()==0) {
		cout<<0<<'\n';
		return ;
	}
	ll ans=n*(n-1)/2+col[cc].size();
	//cout<<"ans="<<ans<<'\n';
	for(int i=col[cc].size()-1;i>=0;i--) {
		int x=col[cc][i];
		//cout<<"x="<<x<<'\n';
		for(int son:g[x]) {
			if(son==fa[x]) continue;
			ll now=cal(son);
		//	cout<<son<<' '<<now<<'\n';
			ans-=now*(now-1)/2;
		}
		ll now=cal(x);
		add(id[x],-now);
		b[x]=now;
	}
	if(c[1]!=cc) {
		ll now=cal(1);
		ans-=now*(now-1)/2;
	}
	for(int x:col[cc]) {
		add(id[x],b[x]);
	}
	cout<<ans<<'\n';
}
bool cmp(int a,int b) {
	return id[a]<id[b];
}
int main()
{
	ios::sync_with_stdio(false);cin.tie(0);
	cin>>n;
	rep(i,1,n) {
		cin>>c[i];
		col[c[i]].pb(i);
	}
	rep(i,1,n-1) {
		int u,v;
		cin>>u>>v;
		g[u].pb(v);
		g[v].pb(u);
	}
	dfs(1,0);
	rep(i,1,n) {
		sort(col[i].begin(),col[i].end(),cmp);
	}
	rep(i,1,n) solve(i);
	return 0;
}


Compiler output:
a.cc: In function 'void add(int, int)':
a.cc:34:44: error: reference to 'data' is ambiguous
   34 |         for(int i=pos;i<maxn;i+=lowbit(i)) data[i]+=x;
      |                                            ^~~~
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:29:4: note:                 'll data [200010]'
   29 | ll data[maxn];
      |    ^~~~
a.cc: In function 'll query(int)':
a.cc:38:44: error: reference to 'data' is ambiguous
   38 |         for(int i=pos;i;i-=lowbit(i)) ans+=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:29:4: note:                 'll data [200010]'
   29 | ll data[maxn];
      |    ^~~~