Submission #12148841


Source Code Expand

Copy
#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];
void 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;
	}
}
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;
		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

Submission Time
Task F - path pass i
User chocorusk
Language C++ (GCC 9.2.1)
Score 600
Code Size 1483 Byte
Status AC
Exec Time 631 ms
Memory 61408 KB

Judge Result

Set Name Sample All
Score / Max Score 0 / 0 600 / 600
Status
AC × 5
AC × 24
Set Name Test Cases
Sample sample_01.txt, sample_02.txt, sample_03.txt, sample_04.txt, sample_05.txt
All line_01.txt, line_3_01.txt, random_01.txt, random_02.txt, random_03.txt, random_100_01.txt, random_100_02.txt, random_10_01.txt, random_10_02.txt, random_1_01.txt, random_1_02.txt, random_2_01.txt, random_2_02.txt, random_2_03.txt, random_3_01.txt, random_3_02.txt, random_3_03.txt, sample_01.txt, sample_02.txt, sample_03.txt, sample_04.txt, sample_05.txt, star_01.txt, star_3_01.txt
Case Name Status Exec Time Memory
line_01.txt AC 535 ms 61408 KB
line_3_01.txt AC 446 ms 61232 KB
random_01.txt AC 631 ms 27176 KB
random_02.txt AC 620 ms 27328 KB
random_03.txt AC 617 ms 27344 KB
random_100_01.txt AC 503 ms 27272 KB
random_100_02.txt AC 504 ms 27296 KB
random_10_01.txt AC 491 ms 27084 KB
random_10_02.txt AC 479 ms 27144 KB
random_1_01.txt AC 469 ms 27116 KB
random_1_02.txt AC 460 ms 27160 KB
random_2_01.txt AC 482 ms 27172 KB
random_2_02.txt AC 471 ms 27168 KB
random_2_03.txt AC 473 ms 27152 KB
random_3_01.txt AC 469 ms 27312 KB
random_3_02.txt AC 468 ms 27192 KB
random_3_03.txt AC 471 ms 27208 KB
sample_01.txt AC 8 ms 8304 KB
sample_02.txt AC 8 ms 8192 KB
sample_03.txt AC 8 ms 8256 KB
sample_04.txt AC 5 ms 8336 KB
sample_05.txt AC 5 ms 8204 KB
star_01.txt AC 479 ms 27416 KB
star_3_01.txt AC 406 ms 27372 KB