forked from SR-Sunny-Raj/Hacktoberfest2021-DSA
-
Notifications
You must be signed in to change notification settings - Fork 0
/
cses tree solution-distinct colours
138 lines (125 loc) · 2.89 KB
/
cses tree solution-distinct colours
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
#include<bits/stdc++.h>
using namespace std;
//problem link
//https://cses.fi/problemset/task/1139
//please merge this as hacktoberfest-accepted label
//this problem uses very good use of pointer .
//hope u like this
//thanks
#define ll long long
#define ld long double
#define rep(i,a,b) for(ll i=a;i<b;i++)
#define repb(i,a,b) for(ll i=a;i>=b;i--)
#define err() cout<<"--------------------------"<<endl;
#define errA(A) for(auto i:A) cout<<i<<" ";cout<<endl;
#define err1(a) cout<<#a<<" "<<a<<endl
#define err2(a,b) cout<<#a<<" "<<a<<" "<<#b<<" "<<b<<endl
#define err3(a,b,c) cout<<#a<<" "<<a<<" "<<#b<<" "<<b<<" "<<#c<<" "<<c<<endl
#define err4(a,b,c,d) cout<<#a<<" "<<a<<" "<<#b<<" "<<b<<" "<<#c<<" "<<c<<" "<<#d<<" "<<d<<endl
#define pb push_back
#define all(A) A.begin(),A.end()
#define allr(A) A.rbegin(),A.rend()
#define ft first
#define sd second
#define pii pair<int,int>
#define pll pair<ll,ll>
#define V vector<ll>
#define S set<ll>
#define VV vector<V>
#define Vpll vector<pll>
#define endl "\n"
ll binaryexp(ll a,ll b,ll i)
{
if(b==0)
return 1LL;
if(b==1)
return a;
ll k=binaryexp(a,b/2,i);
if(b&1)
{
return (((k*k)%i)*a)%i;
}
else
return (k*k)%i;
}
V col(200005);
set<int> *st[200005];
V ans(200005);
void dfs(VV &a,int i,int p)
{
int mxs=-1,mxi=-1;
int n=a.size();
for(auto x:a[i])
{
if(x!=p)
{
dfs(a,x,i);
int jj=st[x]->size();
if(jj>mxs)
{
mxs=st[x]->size();
mxi=x;
}
}
}
// err1(mxi);
if(mxi==-1)
{
st[i]=new set<int>();
st[i]->insert(col[i]);
}
else
{
st[i]=st[mxi];
for(auto x:a[i])
{
if(x!=p && x!=mxi)
{
for(auto kk: *st[x])
st[i]->insert(kk);
}
}
st[i]->insert(col[i]);
}
ans[i]=st[i]->size();
}
int main()
{
ios_base::sync_with_stdio(false);
cin.tie(NULL); cout.tie(NULL);
#ifndef ONLINE_JUDGE
clock_t tStart = clock();
freopen("input.txt","r",stdin);
freopen("output.txt","w",stdout);
#endif
int t;
// cin>>t;
t=1;
while(t--)
{
int n;
cin>>n;
VV a(n);
rep(i,0,n)
{
cin>>col[i];
}
rep(i,0,n-1)
{
int x,y;
cin>>x>>y;
x-=1;
y-=1;
a[x].pb(y);
a[y].pb(x);
}
dfs(a,0,-1);
rep(i,0,n)
{
cout<<ans[i]<<" ";
}
}
#ifndef ONLINE_JUDGE
printf("\nRun Time -> %.10fs\n", (double)(clock()-tStart) / CLOCKS_PER_SEC);
#endif
}