树上启发式合并

Posted 2022-05-28 toot-wjh

模板：

 1 #include<iostream>
 2 #include<cstdio>
 3 #include<cstring>
 4 #include<cmath>
 5 using namespace std;
 6 #define R register
 7 #define ll long long
 8 inline ll read()
 9     ll aa=0;R int bb=1;char cc=getchar();
10     while(cc<‘0‘||cc>‘9‘)
11         if(cc==‘-‘)bb=-1;cc=getchar();
12     while(cc>=‘0‘&&cc<=‘9‘)
13         aa=(aa<<1)+(aa<<3)+(cc^48);cc=getchar();
14     return aa*bb;
15 
16 const int N=1e5+3;
17 struct edge
18     int v,last;
19 ed[N<<1];
20 int first[N],tot;
21 inline void add(int x,int y)
22 
23     ed[++tot].v=y;
24     ed[tot].last=first[x];
25     first[x]=tot;
26 
27 int n,m,c[N],son[N],cnt[N],ans[N],siz[N];
28 void dfsi(int x,int fa)
29 
30     siz[x]=1;
31     for(R int i=first[x],v;i;i=ed[i].last)
32         v=ed[i].v;
33         if(v==fa)continue;
34         dfsi(v,x);
35         siz[x]+=siz[v];
36         if(siz[v]>siz[son[x]])son[x]=v;
37     
38     return;
39 
40 int dfsj(int x,int fa,int bs,int kep)
41 
42     if(kep)
43         for(R int i=first[x],v;i;i=ed[i].last)
44             v=ed[i].v;
45             if(v!=fa&&v!=son[x])
46                 dfsj(v,x,0,1);
47         
48     
49     int res=0;
50     if(son[x])res+=dfsj(son[x],x,1,kep);
51     for(R int i=first[x],v;i;i=ed[i].last)
52         v=ed[i].v;
53         if(v!=fa&&v!=son[x])
54             res+=dfsj(v,x,0,0);
55     
56     if(!cnt[c[x]])res++;
57     cnt[c[x]]++;
58     if(kep)
59         ans[x]=res;
60         if(!bs)memset(cnt,0,sizeof(cnt));
61     
62     return res;
63 
64 int main()
65 
66     n=read();
67     for(R int i=1,x,y;i<n;++i)
68         x=read();y=read();
69         add(x,y);add(y,x);
70     
71     for(R int i=1;i<=n;++i)c[i]=read();
72     dfsi(1,0); dfsj(1,0,1,1);
73     m=read();
74     for(R int i=1,x;i<=m;++i)
75         x=read();
76         printf("%d\n",ans[x]);
77     
78     return 0;
79 

memset奇慢代码 2820ms

 1 // luogu-judger-enable-o2
 2 #include<bits/stdc++.h>
 3 using namespace std;
 4 #define ll int
 5 #define r register 
 6 #define A 1001010
 7 ll head[A],nxt[A],ver[A],size[A],col[A],cnt[A],ans[A],son[A];
 8 ll tot=0,num,sum,nowson,n,m,xx,yy;
 9 inline void add(ll x,ll y)
10     nxt[++tot]=head[x],head[x]=tot,ver[tot]=y;
11 
12 inline ll read()
13     ll f=1,x=0;char c=getchar();
14     while(!isdigit(c))
15         if(c==‘-‘) f=-1;
16         c=getchar();
17     
18     while(isdigit(c))
19         x=(x<<1)+(x<<3)+(c^48),c=getchar();
20     return f*x;
21 
22 void dfs(ll x,ll fa)
23     size[x]=1;
24     for(ll i=head[x];i;i=nxt[i])
25         ll y=ver[i];
26         if(y==fa) continue;
27         dfs(y,x);
28         size[x]+=size[y];
29         if(size[son[x]]<size[y])
30             son[x]=y;
31     
32 
33 void cal(ll x,ll fa,ll val)
34     if(!cnt[col[x]]) ++sum;
35     cnt[col[x]]+=val;
36     for(ll i=head[x];i;i=nxt[i])
37         ll y=ver[i];
38         if(y==fa||y==nowson) continue;
39         cal(y,x,val); 
40     
41 
42 void dsu(ll x,ll fa,bool op)
43     for(ll i=head[x];i;i=nxt[i])
44         ll y=ver[i];
45         if(y==fa||y==son[x])
46             continue;
47         dsu(y,x,0);
48         //从轻儿子出发
49     
50     if(son[x])
51         dsu(son[x],x,1),nowson=son[x];
52     cal(x,fa,1);nowson=0;
53     ans[x]=sum;
54     if(!op)
55         cal(x,fa,-1);
56         sum=0;
57     
58 
59 int main()
60     n=read();
61     for(ll i=1;i<=n-1;i++)
62         xx=read(),yy=read();
63         add(xx,yy),add(yy,xx);
64     
65     for(ll i=1;i<=n;i++)
66         col[i]=read();
67     dfs(1,0);
68     dsu(1,0,1);
69     m=read();
70     for(ll i=1;i<=m;i++)
71         xx=read();
72         printf("%d\n",ans[xx]);
73     
74 

不用memset 381ms