Language: C

#include <bits/stdc++.h>
using namespace std;
#define rep(i, n)  for(long long i=0;i<(long long)(n);i++)
#define REP(i,k,n) for(long long i=k;i<(long long)(n);i++)
#define all(a) a.begin(),a.end()
#define pb emplace_back
#define eb emplace_back
#define lb(v,k) (lower_bound(all(v),k)-v.begin())
#define ub(v,k) (upper_bound(all(v),k)-v.begin())
#define fi first
#define se second
#define pi M_PI
#define PQ(T) priority_queue<T>
#define SPQ(T) priority_queue<T,vector<T>,greater<T>>
#define dame(a) {out(a);return 0;}
#define decimal cout<<fixed<<setprecision(15);
#define dupli(a) erase(unique(all(a)),a.end())
typedef long long ll;
typedef pair<ll,ll> P;
typedef tuple<ll,ll,ll> PP;
typedef tuple<ll,ll,ll,ll> PPP;
typedef multiset<ll> S;
using vi=vector<ll>;
using vvi=vector<vi>;
using vvvi=vector<vvi>;
using vp=vector<P>;
using vvp=vector<vp>;
using vb=vector<bool>;
using vvb=vector<vb>;
const ll inf=1001001001001001001;
const int INF=1001001001;
const int mod=1000000007;
const double eps=1e-10;
template<class T> bool chmin(T&a,T b){if(a>b){a=b;return true;}return false;}
template<class T> bool chmax(T&a,T b){if(a<b){a=b;return true;}return false;}
template<class T> void out(T a){cout<<a<<'\n';}
template<class T> void outp(T a){cout<<'('<<a.fi<<','<<a.se<<')'<<'\n';}
template<class T> void outvp(T v){rep(i,v.size())cout<<'('<<v[i].fi<<','<<v[i].se<<')';cout<<'\n';}
template<class T> void outvvp(T v){rep(i,v.size())outvp(v[i]);}
template<class T> void outv(T v){rep(i,v.size()){if(i)cout<<' ';cout<<v[i];}cout<<'\n';}
template<class T> void outvv(T v){rep(i,v.size())outv(v[i]);}
template<class T> bool isin(T x,T l,T r){return (l)<=(x)&&(x)<=(r);}
template<class T> void yesno(T b){if(b)out("yes");else out("no");}
template<class T> void YesNo(T b){if(b)out("Yes");else out("No");}
template<class T> void YESNO(T b){if(b)out("YES");else out("NO");}
template<class T> void noyes(T b){if(b)out("no");else out("yes");}
template<class T> void NoYes(T b){if(b)out("No");else out("Yes");}
template<class T> void NOYES(T b){if(b)out("NO");else out("YES");}
void outs(ll a,ll b){if(a>=inf-100)out(b);else out(a);}
ll gcd(ll a,ll b){if(b==0)return a;return gcd(b,a%b);}
ll modpow(ll a,ll b){a%=mod;if(b==0)return 1;if(b&1)return a*modpow(a,b-1)%mod;ll k=modpow(a,b/2);return k*k%mod;}
vi dep,par,tour,sz;
vvi g,e;
ll dfs0(int i,int p,int d){
    par[i]=p;
    dep[i]=d;
    tour.pb(i);
    sz[i]=1;
    for(ll x:g[i])if(x!=p){
        sz[i]+=dfs0(x,i,d+1);
        tour.pb(i);
    }
    return sz[i];
}
void init0(ll n){
    g=vvi(n);
    e=vvi(n);
    dep=vi(n);
    par=vi(n);
    sz=vi(n);
    rep(i,n-1){
        ll a,b;cin>>a>>b;a--;b--;
        g[a].pb(b);
        g[b].pb(a);
    }
    dfs0(0,-1,0);
    rep(i,tour.size())e[tour[i]].pb(i);
}
vi seg;
ll N;
void init(ll n){
    N=1;
    while(N<n)N*=2;
    seg=vi(N*2-1);
}
void update(ll i,ll x){
    i=N+i-1;
    seg[i]=x;
    while(i>0){
        i=(i-1)/2;
        seg[i]=seg[i*2+1]+seg[i*2+2];
    }
}
ll getsum(ll a,ll b,ll k,ll l,ll r){
    if(r<=a||b<=l)return 0;
    if(a<=l&&r<=b)return seg[k];
    ll c1=getsum(a,b,k*2+1,l,(l+r)/2);
    ll c2=getsum(a,b,k*2+2,(l+r)/2,r);
    return c1+c2;
}
int main(){
    ll n;cin>>n;
    vi color(n);
    rep(i,n){
        cin>>color[i];color[i]--;
    }
    init0(n);
    vvp s(n);
    rep(i,n)s[color[i]].eb(dep[i],i);
    rep(i,n){
        sort(all(s[i]));reverse(all(s[i]));
    }
    vi ans(n);
    init(tour.size());
    rep(tt,n){
        vi change;
        for(auto x:s[tt]){
            ll i=x.se;
            for(ll to:g[i])if(to!=par[i]){
                ll k=sz[to]-getsum(e[to][0],e[to].back()+1,0,0,N);
                ans[tt]+=k*(k+1)/2;
            }
            update(e[i][0],sz[i]-getsum(e[i][0],e[i].back()+1,0,0,N));
            change.pb(e[i][0]);
        }
        ll k=sz[0]-getsum(e[0][0],e[0].back()+1,0,0,N);
        ans[tt]+=k*(k+1)/2;
        for(ll x:change)update(x,0);
    }
    for(ll x:ans)out(n*(n+1)/2-x);
}

Compiler output:
main.c:1:10: fatal error: bits/stdc++.h: No such file or directory
    1 | #include <bits/stdc++.h>
      |          ^~~~~~~~~~~~~~~
compilation terminated.