题意:一棵树上有n(n<=50000)个结点,结点有k(k<=10)种颜色,问树上总共有多少条包含所有颜色的路径。
我最初的想法是树形状压dp,设dp[u][S]为以结点u为根的包含颜色集合为S的路径条数,然后FWT(应该叫FMT?)搞一下就行了,复杂度$O(nk2^k)$。奈何内存太大,妥妥地MLE...
看到网上大部分的解法都是点分治,我不禁联想到之前学过的树上任意两点距离的求法(点分治+FFT),心想,这道题用点分治+FWT是不是也能过?于是比着葫芦画瓢写出了这样一段又臭又长的代码:
1 #include<bits/stdc++.h> 2 using namespace std; 3 typedef long long ll; 4 typedef double db; 5 const int N=5e4+10,inf=0x3f3f3f3f; 6 int n,k,a[N],hd[N],ne,vis[N],K,siz[N],tot,rt,mx; 7 ll dp[1<<10],ans; 8 struct E {int v,nxt;} e[N<<1]; 9 void addedge(int u,int v) {e[ne]= {v,hd[u]},hd[u]=ne++;} 10 void FWT(ll* a,int n,int f) { 11 for(int k=1; k<n; k<<=1) 12 for(int i=0; i<n; i+=k<<1) 13 for(int j=i; j<i+k; ++j) 14 a[j+k]+=f==1?a[j]:-a[j]; 15 } 16 void getroot(int u,int fa) { 17 siz[u]=1; 18 int sz=0; 19 for(int i=hd[u]; ~i; i=e[i].nxt) { 20 int v=e[i].v; 21 if(vis[v]||v==fa)continue; 22 getroot(v,u); 23 siz[u]+=siz[v]; 24 sz=max(sz,siz[v]); 25 } 26 sz=max(sz,tot-siz[u]); 27 if(sz<mx)mx=sz,rt=u; 28 } 29 void dfs(int u,int fa,int S) { 30 dp[S]++; 31 for(int i=hd[u]; ~i; i=e[i].nxt) { 32 int v=e[i].v; 33 if(vis[v]||v==fa)continue; 34 dfs(v,u,S|a[v]); 35 } 36 } 37 void cal(int u,int ba,int f) { 38 for(int i=0; i<=K; ++i)dp[i]=0; 39 dfs(u,-1,a[u]|ba); 40 FWT(dp,K+1,1); 41 for(int i=0; i<=K; ++i)dp[i]*=dp[i]; 42 FWT(dp,K+1,-1); 43 ans+=dp[K]*f; 44 } 45 void solve(int u) { 46 mx=inf,getroot(u,-1),u=rt,cal(u,0,1),vis[u]=1; 47 for(int i=hd[u]; ~i; i=e[i].nxt) { 48 int v=e[i].v; 49 if(!vis[v])tot=siz[v],cal(v,a[u],-1),solve(v); 50 } 51 } 52 ll treepartion() { 53 ans=0,tot=n; 54 solve(1); 55 return ans; 56 } 57 int main() { 58 while(scanf("%d%d",&n,&k)==2) { 59 memset(hd,-1,sizeof hd),ne=0; 60 memset(vis,0,sizeof vis); 61 K=(1<<k)-1; 62 for(int i=1; i<=n; ++i)scanf("%d",&a[i]),a[i]=1<<(a[i]-1); 63 for(int i=1; i<n; ++i) { 64 int u,v; 65 scanf("%d%d",&u,&v); 66 addedge(u,v); 67 addedge(v,u); 68 } 69 printf("%lld\n",treepartion()); 70 } 71 return 0; 72 }
虽然成功地AC了,但是仔细一想:不对啊,这道题FWT的复杂度和子树的大小不是线性相关的啊!所以这样一来,总的复杂度成了$O(nk2^klogn)$,反而增大了。
也就是说,这道题用点分治的作用仅仅是减少了内存的开销,复杂度非但没有减少,反而还多了个logn!
当然除了点分治,这道题还有其他的优化方法,比如sclbgw7大佬利用树链剖分的思想将内存优化到了$O(2^klogn)$,时间复杂度仍为$O(nk2^k)$。
还有Menhera大佬利用基的FMT性质将时间复杂度优化到$O(n2^k)$的做法,看样子有点像是容斥。我个人更倾向于这一种,于是在这个思想的基础上进一步地分析:
题目要求的是包含所有颜色的路径条数。如果包含某个元素集合的路径不太好求,那么不包含某个元素集合的呢?只要把属于这个集合的结点都染成白色,不属于的都染成黑色,则问题就转化成了求一棵树上包含的所有点都是黑色的路径条数,直接dp求一下就行了。于是我们可以利用容斥原理,用所有的路径数减去不包含1个元素集合的路径数,再加上不包含2个元素集合的路径数,再减去不包含3个...就得到了答案。
1 #include<bits/stdc++.h> 2 using namespace std; 3 typedef long long ll; 4 const int N=5e4+10,inf=0x3f3f3f3f; 5 int n,k,a[N],siz[N],hd[N],ne,dp[N],ppc[1<<10]; 6 ll ans; 7 struct E {int v,nxt;} e[N<<1]; 8 void addedge(int u,int v) {e[ne]= {v,hd[u]},hd[u]=ne++;} 9 void dfs(int u,int fa,int f) { 10 for(int i=hd[u]; ~i; i=e[i].nxt)if(e[i].v!=fa)dfs(e[i].v,u,f); 11 if(!siz[u])return; 12 ans+=f; 13 for(int i=hd[u]; ~i; i=e[i].nxt)if(e[i].v!=fa)siz[u]+=siz[e[i].v]; 14 for(int i=hd[u]; ~i; i=e[i].nxt)if(e[i].v!=fa) 15 ans+=(ll)siz[e[i].v]*(siz[u]-siz[e[i].v]+1)*f; 16 } 17 ll solve() { 18 ans=0; 19 for(int S=(1<<k)-1; S; --S) { 20 int f=(k-ppc[S])&1?-1:1; 21 for(int i=1; i<=n; ++i)siz[i]=S>>a[i]&1; 22 dfs(1,-1,f); 23 } 24 return ans; 25 } 26 int main() { 27 ppc[0]=0; 28 for(int i=1; i<(1<<10); ++i)ppc[i]=ppc[i>>1]+(i&1); 29 while(scanf("%d%d",&n,&k)==2) { 30 memset(hd,-1,sizeof hd),ne=0; 31 for(int i=1; i<=n; ++i)scanf("%d",&a[i]),a[i]--; 32 for(int i=1; i<n; ++i) { 33 int u,v; 34 scanf("%d%d",&u,&v); 35 addedge(u,v); 36 addedge(v,u); 37 } 38 printf("%lld\n",solve()); 39 } 40 return 0; 41 }
这种方法的复杂度为什么会比FWT少了k呢?这个k哪里去了呢?我想大概是在FWT的过程中把所有集合的dp值都求出来了,而我们只需要求全集的dp值,因此多做了许多无用功。
(ps:由于题目数据的限制,用map优化的点分治可能会更快一些)
HDU - 5977 Garden of Eden (容斥)
原文:https://www.cnblogs.com/asdfsag/p/10557845.html