This submission is migrated from previous version of oj.uz, which used different machine for grading. This submission may have different result if resubmitted.
#include <bits/stdc++.h>
#ifdef NON_SUBMIT
#define TEST(n) (n)
#define tout cerr
#else
#define TEST(n) ((void)0)
#define tout cin
#endif
using namespace std;
vector<int> adj[200000], Q[400000], tree[400000];
int res, N, node_cnt, st_cnt, STP[200000], num[200000], num_rev[200000], ST[400000][19], SZ[400000], depth[200000], parent[200000], P[400000][18], dist[400000], sdist[400000], ans[200000], C[200000], CC[400000], cnt[200001];
void dfs(int c)
{
STP[c]=st_cnt;
ST[st_cnt][0]=num[c]=node_cnt++;
num_rev[num[c]]=c;
for(int i=1;i<19;i++) ST[st_cnt][i]=min(ST[st_cnt][i-1],ST[max(st_cnt-(1<<i-1),0)][i-1]);
st_cnt++;
for(auto n: adj[c]) if(num[n]==-1) {
parent[n]=c;
depth[n]=depth[c]+1;
dfs(n);
ST[st_cnt][0]=num[c];
for(int i=1;i<19;i++) ST[st_cnt][i]=min(ST[st_cnt][i-1],ST[max(st_cnt-(1<<i-1),0)][i-1]);
st_cnt++;
}
}
int LCA(int a, int b)
{
int d;
if(STP[a]>STP[b]) swap(a,b);
d=SZ[STP[b]-STP[a]+1];
a=STP[a]; b=STP[b];
return num_rev[min(ST[a+(1<<d)-1][d],ST[b][d])];
}
int get_dist(int a, int b)
{
b=b<N ? b:parent[b-N];
return depth[a]+depth[b]-2*depth[LCA(a,b)];
}
void solve(int c)
{
for(auto n: adj[c]) if(num[n]>num[c]) {
solve(n);
if(dist[c]<dist[n]+1) {
sdist[c]=dist[c];
dist[c]=dist[n]+1;
}
else if(sdist[c]<dist[n]+1) sdist[c]=dist[n]+1;
}
P[c][0]=c;
for(auto n: adj[c]) if(num[n]>num[c] && sdist[c]<dist[n]+1) {
int np=n;
for(int i=17;i>=0;i--) if(get_dist(c,P[np][i])<=sdist[c]) np=P[np][i];
if(get_dist(c,np)>sdist[c]) P[c][0]=np;
else if(get_dist(c,P[np][0])>sdist[c]) P[c][0]=P[np][0];
}
for(int j=1;j<18;j++) P[c][j]=P[P[c][j-1]][j-1];
}
void solve2(int c)
{
int m=N+c, sm=N, tm=N;
for(int j=1;j<18;j++) P[N+c][j]=P[P[N+c][j-1]][j-1];
for(auto n: adj[c]) if(num[n]>num[c]) {
if(dist[m]<dist[n]) {
tm=sm;
sm=m;
m=n;
}
else if(dist[sm]<dist[n]) {
tm=sm;
sm=n;
}
else if(dist[tm]<dist[n]) tm=n;
}
for(auto n: adj[c]) if(num[n]>num[c]) {
int np;
dist[N+n]=dist[m]==dist[n] ? dist[sm]+1:dist[m]+1;
if(m==n) {
np=sm;
for(int i=17;i>=0;i--) if(get_dist(c,P[np][i])<=dist[tm]+1) np=P[np][i];
if(get_dist(c,np)>dist[tm]+1) P[N+n][0]=np;
else if(get_dist(c,P[np][0])>dist[tm]+1) P[N+n][0]=P[np][0];
}
else if(sm==n) {
np=m;
for(int i=17;i>=0;i--) if(get_dist(c,P[np][i])<=dist[tm]+1) np=P[np][i];
if(get_dist(c,np)>dist[tm]+1) P[N+n][0]=np;
else if(get_dist(c,P[np][0])>dist[tm]+1) P[N+n][0]=P[np][0];
}
else {
np=m;
for(int i=17;i>=0;i--) if(get_dist(c,P[np][i])<=dist[sm]+1) np=P[np][i];
if(get_dist(c,np)>dist[sm]+1) P[N+n][0]=np;
else if(get_dist(c,P[np][0])>dist[sm]+1) P[N+n][0]=P[np][0];
}
solve2(n);
}
}
void dfs2(int c)
{
CC[c]=2;
if(++cnt[C[c<N ? c:parent[c-N]]]==1) res++;
for(auto i: Q[c]) ans[i]=res;
for(auto n: tree[c]) if(CC[n]==0) dfs2(n);
if(--cnt[C[c<N ? c:parent[c-N]]]==0) res--;
}
int main()
{
ios::sync_with_stdio(false);
cin.tie(NULL); cout.tie(NULL);
TEST(freopen("input.txt","r",stdin));
TEST(freopen("output.txt","w",stdout));
TEST(freopen("debug.txt","w",stderr));
int M;
cin>>N>>M;
memset(P,-1,sizeof(P));
for(int i=2;i<400000;i<<=1) SZ[i+1]=1;
for(int i=2;i<400000;i++) SZ[i]+=SZ[i-1];
dist[P[N][0]=N]=-1;
for(int i=1;i<N;i++) {
int a, b;
cin>>a>>b;
adj[--a].push_back(--b);
adj[b].push_back(a);
P[N+i][0]=N+i;
num[i]=-1;
}
dfs(0); solve(0); solve2(0);
for(int i=0;i<N;i++) {
cin>>C[i];
if(dist[i]>dist[N+i]+1) {
int c=i;
for(int j=17;j>=0;j--) if(get_dist(i,P[c][j])<=dist[N+i]+1) c=P[c][j];
if(get_dist(i,c)>dist[N+i]+1) Q[c].push_back(i);
else if(get_dist(i,P[c][0])>dist[N+i]+1) Q[P[c][0]].push_back(i);
}
else if(dist[i]<dist[N+i]+1) {
int c=N+i;
for(int j=17;j>=0;j--) if(get_dist(i,P[c][j])<=dist[i]) c=P[c][j];
if(get_dist(i,c)>dist[i]) Q[c].push_back(i);
else if(get_dist(i,P[c][0])>dist[i]) Q[P[c][0]].push_back(i);
}
}
for(int i=0;i<2*N;i++) tree[P[i][0]].push_back(i);
CC[N]=2;
for(int i=0;i<2*N;i++) if(CC[i]==0) {
int r=P[P[i][17]][17];
for(;CC[r]==0;r=P[r][0]) {
CC[r]=1;
if(++cnt[C[r<N ? r:parent[r-N]]]==1) res++;
}
for(;CC[r]==1;r=P[r][0]) dfs2(r);
for(;CC[r]==2;r=P[r][0]) {
CC[r]=3;
if(--cnt[C[r<N ? r:parent[r-N]]]==0) res--;
}
}
for(int i=0;i<N;i++) cout<<ans[i]<<'\n';
return 0;
}
Compilation message (stderr)
joi2019_ho_t5.cpp: In function 'void dfs(int)':
joi2019_ho_t5.cpp:21:77: warning: suggest parentheses around '-' inside '<<' [-Wparentheses]
for(int i=1;i<19;i++) ST[st_cnt][i]=min(ST[st_cnt][i-1],ST[max(st_cnt-(1<<i-1),0)][i-1]);
~^~
joi2019_ho_t5.cpp:28:78: warning: suggest parentheses around '-' inside '<<' [-Wparentheses]
for(int i=1;i<19;i++) ST[st_cnt][i]=min(ST[st_cnt][i-1],ST[max(st_cnt-(1<<i-1),0)][i-1]);
~^~
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |