HDU5909Tree Cutting（FWT）

Posted 2020-11-07 小蒟蒻yyb的博客

【HDU5909】Tree Cutting（FWT）

题面

vjudge
题目大意：
给你一棵\(n\)个节点的树，每个节点都有一个小于\(m\)的权值
定义一棵子树的权值为所有节点的异或和，问权值为\(0..m-1\)的所有子树的个数

题解

考虑\(dp\)
设\(f[i][j]\)表示以\(i\)为根节点的子树中，异或和为\(j\)的子树的个数
很显然，每次合并就是两个\(dp\)值做\(xor\)卷积
那么直接用\(FWT\)优化就行了

#include<iostream>
#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<cmath>
#include<algorithm>
#include<set>
#include<map>
#include<vector>
#include<queue>
using namespace std;
#define ll long long
#define RG register
#define MAX 1111
#define MOD (1000000007)
#define inv2 (500000004)
inline int read()
{
    RG int x=0,t=1;RG char ch=getchar();
    while((ch<'0'||ch>'9')&&ch!='-')ch=getchar();
    if(ch=='-')t=-1,ch=getchar();
    while(ch<='9'&&ch>='0')x=x*10+ch-48,ch=getchar();
    return x*t;
}
struct Line{int v,next;}e[MAX<<1];
int h[MAX],cnt;
inline void Add(int u,int v){e[cnt]=(Line){v,h[u]};h[u]=cnt++;}
int n,m,N;
void FWT(int *P,int opt)
{
    for(int i=2;i<=N;i<<=1)
        for(int p=i>>1,j=0;j<N;j+=i)
            for(int k=j;k<j+p;++k)
            {
                int x=P[k],y=P[k+p];
                P[k]=(x+y)%MOD;P[k+p]=(x-y+MOD)%MOD;
                if(opt==-1)P[k]=1ll*P[k]*inv2%MOD,P[k+p]=1ll*P[k+p]*inv2%MOD; 
            }
}
int f[MAX][1<<10],ans[MAX];
void dfs(int u,int ff)
{
    FWT(f[u],1);
    for(int i=h[u];i;i=e[i].next)
    {
        int v=e[i].v;if(v==ff)continue;
        dfs(v,u);
        for(int j=0;j<N;++j)f[u][j]=1ll*f[u][j]*f[v][j]%MOD;
    }
    FWT(f[u],-1);f[u][0]=(f[u][0]+1)%MOD;FWT(f[u],1);
}
int main()
{
    int T=read();
    while(T--)
    {
        n=read();m=read();cnt=1;N=m;
        memset(f,0,sizeof(f));memset(h,0,sizeof(h));memset(ans,0,sizeof(ans));
        for(int i=1;i<=n;++i)f[i][read()]++;
        for(int i=1;i<n;++i)
        {
            int u=read(),v=read();
            Add(u,v);Add(v,u);
        }
        dfs(1,0);
        for(int i=1;i<=n;++i)FWT(f[i],-1);
        for(int i=1;i<=n;++i)f[i][0]=(f[i][0]-1+MOD)%MOD;
        for(int i=1;i<=n;++i)
            for(int j=0;j<m;++j)
                ans[j]=(ans[j]+f[i][j])%MOD;
        for(int i=0;i<m;++i,(i==m)?putchar('\n'):putchar(' '))printf("%d",ans[i]);
    }
}