Submission #12148782

---

Source Code Expand

#include <cstdio>
#include <cstring>
#include <iostream>
#include <string>
#include <cmath>
#include <bitset>
#include <vector>
#include <map>
#include <set>
#include <queue>
#include <deque>
#include <algorithm>
#include <complex>
#include <unordered_map>
#include <unordered_set>
#include <random>
#include <cassert>
#include <fstream>
#include <utility>
#include <functional>
#include <time.h>
#include <stack>
#include <array>
#include <list>
#define popcount __builtin_popcount
using namespace std;
typedef long long int ll;
typedef pair<int, int> P;
vector<int> g[200020];
int n;
int c[200020], sz[200020];
ll s[200020];
map<int, int> merge(map<int, int> a, map<int, int> b){
	if(a.size()<b.size()) swap(a, b);
	for(auto p:b){
		a[p.first]+=p.second;
	}
	return a;
}
void dfs(int x, int p, map<int, int> &mp){
	sz[x]=1;
	for(auto y:g[x]){
		if(y==p) continue;
		map<int, int> mpy;
		dfs(y, x, mpy);
		sz[x]+=sz[y];
		ll c1=sz[y];
		if(mpy.find(c[x])!=mpy.end()){
			c1-=mpy[c[x]];
		}
		s[c[x]]+=c1*(c1+1)/2;
		mp=merge(mp, mpy);
	}
	mp[c[x]]=sz[x];
}
int main()
{
	cin>>n;
	for(int i=0; i<n; i++){
		cin>>c[i]; c[i]--;
	}
	for(int i=0; i<n-1; i++){
		int a, b; cin>>a>>b; a--; b--;
		g[a].push_back(b);
		g[b].push_back(a);
	}
	map<int, int> mp;
	dfs(0, -1, mp);
	for(int i=0; i<n; i++){
		if(i==c[0]) continue;
		ll c1=n-mp[i];
		s[i]+=c1*(c1+1)/2;
	}
	for(int i=0; i<n; i++) cout<<(ll)n*(n+1)/2-s[i]<<endl;
	return 0;
}

Submission Info

Judge Result