Submission #12148782


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];
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

Submission Time
Task F - path pass i
User chocorusk
Language C++ (GCC 9.2.1)
Score 0
Code Size 1505 Byte
Status TLE
Exec Time 2207 ms
Memory 86424 KB

Judge Result

Set Name Sample All
Score / Max Score 0 / 0 0 / 600
Status
AC × 5
AC × 19
TLE × 5
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 TLE 2207 ms 77632 KB
line_3_01.txt AC 474 ms 86424 KB
random_01.txt TLE 2206 ms 25460 KB
random_02.txt TLE 2206 ms 23904 KB
random_03.txt TLE 2206 ms 28916 KB
random_100_01.txt AC 632 ms 27564 KB
random_100_02.txt AC 638 ms 27288 KB
random_10_01.txt AC 526 ms 27220 KB
random_10_02.txt AC 529 ms 27344 KB
random_1_01.txt AC 480 ms 27328 KB
random_1_02.txt AC 486 ms 27308 KB
random_2_01.txt AC 494 ms 27320 KB
random_2_02.txt AC 500 ms 27360 KB
random_2_03.txt AC 498 ms 27300 KB
random_3_01.txt AC 492 ms 27504 KB
random_3_02.txt AC 498 ms 27448 KB
random_3_03.txt AC 494 ms 27296 KB
sample_01.txt AC 10 ms 8336 KB
sample_02.txt AC 8 ms 8280 KB
sample_03.txt AC 8 ms 8256 KB
sample_04.txt AC 8 ms 8308 KB
sample_05.txt AC 7 ms 8296 KB
star_01.txt TLE 2206 ms 16812 KB
star_3_01.txt AC 431 ms 27312 KB