HDU - 5977 Garden of Eden (树形dp+容斥)

题意：一棵树上有n(n<=50000)个结点，结点有k(k<=10)种颜色，问树上总共有多少条包含所有颜色的路径。

我最初的想法是树形状压dp，设dp[u][S]为以结点u为根的包含颜色集合为S的路径条数，然后FWT(应该叫FMT？)搞一下就行了，复杂度$O(nk2^k)$。奈何内存太大，妥妥地MLE...

看到网上大部分的解法都是点分治，我不禁联想到之前学过的树上任意两点距离的求法（点分治+FFT），心想，这道题用点分治+FWT是不是也能过？于是比着葫芦画瓢写出了这样一段又臭又长的代码：

 #include<bits/stdc++.h>
 using namespace std;
 typedef long long ll;
 typedef double db;
 const int N=5e4+,inf=0x3f3f3f3f;
 int n,k,a[N],hd[N],ne,vis[N],K,siz[N],tot,rt,mx;
 ll dp[<<],ans;
 struct E {int v,nxt;} e[N<<];
 void addedge(int u,int v) {e[ne]= {v,hd[u]},hd[u]=ne++;}
 void FWT(ll* a,int n,int f) {
     for(int k=; k<n; k<<=)
         for(int i=; i<n; i+=k<<)
             for(int j=i; j<i+k; ++j)
                 a[j+k]+=f==?a[j]:-a[j];
 }
 void getroot(int u,int fa) {
     siz[u]=;
     int sz=;
     for(int i=hd[u]; ~i; i=e[i].nxt) {
         int v=e[i].v;
         if(vis[v]||v==fa)continue;
         getroot(v,u);
         siz[u]+=siz[v];
         sz=max(sz,siz[v]);
     }
     sz=max(sz,tot-siz[u]);
     if(sz<mx)mx=sz,rt=u;
 }
 void dfs(int u,int fa,int S) {
     dp[S]++;
     for(int i=hd[u]; ~i; i=e[i].nxt) {
         int v=e[i].v;
         if(vis[v]||v==fa)continue;
         dfs(v,u,S|a[v]);
     }
 }
 void cal(int u,int ba,int f) {
     for(int i=; i<=K; ++i)dp[i]=;
     dfs(u,-,a[u]|ba);
     FWT(dp,K+,);
     for(int i=; i<=K; ++i)dp[i]*=dp[i];
     FWT(dp,K+,-);
     ans+=dp[K]*f;
 }
 void solve(int u) {
     mx=inf,getroot(u,-),u=rt,cal(u,,),vis[u]=;
     for(int i=hd[u]; ~i; i=e[i].nxt) {
         int v=e[i].v;
         if(!vis[v])tot=siz[v],cal(v,a[u],-),solve(v);
     }
 }
 ll treepartion() {
     ans=,tot=n;
     solve();
     return ans;
 }
 int main() {
     while(scanf("%d%d",&n,&k)==) {
         memset(hd,-,sizeof hd),ne=;
         memset(vis,,sizeof vis);
         K=(<<k)-;
         for(int i=; i<=n; ++i)scanf("%d",&a[i]),a[i]=<<(a[i]-);
         for(int i=; i<n; ++i) {
             int u,v;
             scanf("%d%d",&u,&v);
             addedge(u,v);
             addedge(v,u);
         }
         printf("%lld\n",treepartion());
     }
     return ;
 }

虽然成功地AC了，但是仔细一想：不对啊，这道题FWT的复杂度和子树的大小不是线性相关的啊！所以这样一来，总的复杂度成了$O(nk2^klogn)$，反而增大了。

也就是说，这道题用点分治的作用仅仅是减少了内存的开销，复杂度非但没有减少，反而还多了个logn！

当然除了点分治，这道题还有其他的优化方法，比如sclbgw7大佬利用树链剖分的思想将内存优化到了$O(2^klogn)$，时间复杂度仍为$O(nk2^k)$。

树剖做法：

 #include<bits/stdc++.h>
 using namespace std;
 typedef long long ll;
 const int N=5e4+,inf=0x3f3f3f3f;
 int n,k,a[N],hd[N],ne,S,fa[N],son[N],siz[N],tot;
 ll dp[][<<],b[<<],A[N],B[N],ans;
 struct E {int v,nxt;} e[N<<];
 void addedge(int u,int v) {e[ne]= {v,hd[u]},hd[u]=ne++;}
 void FWT(ll* a,int n,int f) {
     for(int k=; k<n; k<<=)
         for(int i=; i<n; i+=k<<)
             for(int j=i; j<i+k; ++j) {
                 ll x=a[j],y=a[j+k];
                 a[j+k]=f==?y+x:y-x;
             }
 }
 void mul(ll* a,ll* b,ll* c,int n) {
     for(int i=; i<n; ++i)A[i]=a[i],B[i]=b[i];
     FWT(A,n,),FWT(B,n,);
     for(int i=; i<n; ++i)c[i]=A[i]*B[i];
     FWT(c,n,-);
 }
 int newnode() {
     int u=tot++;
     for(int i=; i<(<<k); ++i)dp[u][i]=;
     return u;
 }
 void dfs1(int u,int f) {
     fa[u]=f,son[u]=,siz[u]=;
     for(int i=hd[u]; ~i; i=e[i].nxt) {
         int v=e[i].v;
         if(v==fa[u])continue;
         dfs1(v,u),siz[u]+=siz[v];
         if(siz[v]>siz[son[u]])son[u]=v;
     }
 }
 void dfs2(int u,int w) {
     if(son[u])dfs2(son[u],w);
     for(int i=; i<(<<k); ++i)b[i]=;
     b[<<a[u]]=;
     if((<<a[u]==S))ans++;
     for(int i=; i<(<<k); ++i)if((i|(<<a[u]))==S)ans+=dp[w][i]*;
     for(int i=; i<(<<k); ++i)b[i|(<<a[u])]+=dp[w][i];
     for(int i=; i<(<<k); ++i)dp[w][i]=b[i];
     for(int i=hd[u]; ~i; i=e[i].nxt) {
         int v=e[i].v;
         if(v==fa[u]||v==son[u])continue;
         int wv=newnode();
         dfs2(v,wv),tot--;
         mul(dp[w],dp[wv],b,<<k);
         ans+=b[S]*;
         for(int i=; i<(<<k); ++i)dp[w][i|(<<a[u])]+=dp[wv][i];
     }
 }
 int main() {
     while(scanf("%d%d",&n,&k)==) {
         memset(hd,-,sizeof hd),ne=;
         S=(<<k)-;
         for(int i=; i<=n; ++i)scanf("%d",&a[i]),a[i]--;
         for(int i=; i<n; ++i) {
             int u,v;
             scanf("%d%d",&u,&v);
             addedge(u,v);
             addedge(v,u);
         }
         tot=ans=;
         dfs1(,-),dfs2(,newnode());
         printf("%lld\n",ans);
     }
     return ;
 }

还有Menhera大佬利用基的FMT性质将时间复杂度优化到$O(n2^k)$的做法，看样子有点像是容斥。我个人更倾向于这一种，于是在这个思想的基础上进一步地分析：

题目要求的是包含所有颜色的路径条数。如果包含某个元素集合的路径不太好求，那么不包含某个元素集合的呢？只要把属于这个集合的结点都染成白色，不属于的都染成黑色，则问题就转化成了求一棵树上包含的所有点都是黑色的路径条数，直接dp求一下就行了。于是我们可以利用容斥原理，用所有的路径数减去不包含1个元素集合的路径数，再加上不包含2个元素集合的路径数，再减去不包含3个...就得到了答案。

 #include<bits/stdc++.h>
 using namespace std;
 typedef long long ll;
 const int N=5e4+,inf=0x3f3f3f3f;
 int n,k,a[N],siz[N],hd[N],ne,ppc[<<];
 ll ans;
 struct E {int v,nxt;} e[N<<];
 void addedge(int u,int v) {e[ne]= {v,hd[u]},hd[u]=ne++;}
 void dfs(int u,int fa,int f) {
     for(int i=hd[u]; ~i; i=e[i].nxt)if(e[i].v!=fa)dfs(e[i].v,u,f);
     if(!siz[u])return;
     ans+=f;
     for(int i=hd[u]; ~i; i=e[i].nxt)if(e[i].v!=fa) {
             int v=e[i].v;
             ans+=(ll)siz[v]*siz[u]**f;
             siz[u]+=siz[v];
         }
 }
 ll solve() {
     ans=;
     for(int S=(<<k)-; S; --S) {
         int f=(k-ppc[S])&?-:;
         for(int i=; i<=n; ++i)siz[i]=S>>a[i]&;
         dfs(,-,f);
     }
     return ans;
 }
 int main() {
     ppc[]=;
     for(int i=; i<(<<); ++i)ppc[i]=ppc[i>>]+(i&);
     while(scanf("%d%d",&n,&k)==) {
         memset(hd,-,sizeof hd),ne=;
         for(int i=; i<=n; ++i)scanf("%d",&a[i]),a[i]--;
         for(int i=; i<n; ++i) {
             int u,v;
             scanf("%d%d",&u,&v);
             addedge(u,v);
             addedge(v,u);
         }
         printf("%lld\n",solve());
     }
     return ;
 }

这种方法的复杂度为什么会比FWT少了k呢？这个k哪里去了呢？我想大概是在FWT的过程中把所有集合的dp值都求出来了，而我们只需要求全集的dp值，因此多做了许多无用功。

(ps：由于题目数据的限制，用map优化的点分治可能会更快一些)