【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]);

	}

}