题目描述
When God made the first man, he put him on a beautiful garden, the Garden of Eden. Here Adam lived with all animals. God gave Adam eternal life. But Adam was lonely in the garden, so God made Eve. When Adam was asleep one night, God took a rib from him and made Eve beside him. God said to them, “here in the Garden, you can do everything, but you cannot eat apples from the tree of knowledge.”
One day, Satan came to the garden. He changed into a snake and went to live in the tree of knowledge. When Eve came near the tree someday, the snake called her. He gave her an apple and persuaded her to eat it. Eve took a bite, and then she took the apple to Adam. And Adam ate it, too. Finally, they were driven out by God and began a hard journey of life.
The above is the story we are familiar with. But we imagine that Satan love knowledge more than doing bad things. In Garden of Eden, the tree of knowledge has n apples, and there are k varieties of apples on the tree. Satan wants to eat all kinds of apple to gets all kinds of knowledge.So he chooses a starting point in the tree,and starts walking along the edges of tree,and finally stops at a point in the tree(starting point and end point may be same).The same point can only be passed once.He wants to know how many different kinds of schemes he can choose to eat all kinds of apple. Two schemes are different when their starting points are different or ending points are different.
Input
There are several cases.Process till end of input.
For each case, the first line contains two integers n and k, denoting the number of apples on the tree and number of kinds of apple on the tree respectively.
The second line contains n integers meaning the type of the i-th apple. Types are represented by integers between 1 and k .
Each of the following n-1 lines contains two integers u and v,meaning there is one edge between u and v.1≤n≤50000, 1≤k≤10
Output
For each case output your answer on a single line.
Sample Input
3 2
1 2 2
1 2
1 3
Sample Output
6
题目分析
鸡冻，终于自己完全独立地完成了一道树分治的题目啦。虽然还是调了三四个小时。。。
题目要求的是访问所有颜色的树上节点的方案数，而且必须是最短路径（不能重复访问节点保证的），因此我们联想到树分治。
不同于常见的树分治都是求树上路径如何，这里要求的是路径上点的状态要满足的条件。为了将问题转化我们将点的状态放在路径上，然后转化成树分治进行解决。现在问题转化成了如何表示树的状态。发现节点的颜色很少，我们可以用状态压缩，将颜色转化成二进制的位数，颜色为x，则对应的状态就是1<<(x-1)，转化成数字一来方便表示，二来方便进行操作。
我们成功得到重心到所有点的状态后如何求满足条件的点对数呢？所谓满足条件，就是两个点所表示的状态进行位操作的或运算以后等于(1<<k)-1。
那么我们如何快速得到任意两个点或运算以后等于(1<<k)-1的状态数呢？显然不能直接进行暴力。我们需要联想到FWT。
将每个状态数的点数存放在状态数组中，然后再对这个数组进行FWT，则Cnt[(1<<k)-1]中保存的就是答案。
其他的操作就是树分治的基本操作。
经验总结
虽然很快想到了思路，并且写出了代码，而且也可以通过样例，但是不知道为什么一直WA，而且很不自信，觉得自己写的思路可能有问题。在网上看其他人的代码他们好像都没有用到FWT，好像都是树形DP什么的，那个我还不会。但是又觉得自己的应该没有什么大问题。就很绝望。
睡了一晚起来再慢慢调试，一步一步地看，终于发现自己在一个很细小的地方犯了一个错误，虽然不知道为什么自己捏的各种各样的样例都能过，但是这个错误很根本。改了以后就过啦。
人呐，还是要自信一点，不要遇到错误就着急看别人怎么样。
AC代码
#include<iostream>
#include<cstring>
#include<cstdio>
#include<climits>
#include<algorithm>
#include<ctime>
#include<cstdlib>
#include<queue>
#include<set>
#include<map>
#include<cmath>
using namespace std;
typedef long long ll;
const int MAXN=5e4+5;
struct edge
{
int to,last;
}Edge[MAXN<<2]; int Last[MAXN],tot;
int n,kk,SonNum[MAXN],MaxNum[MAXN],Vis[MAXN],Record[MAXN],Color[MAXN];
int root,rootx,dlen,ss,len;
ll ans,Cnt[(1<<11)+100];
int Power[15];
void FWT(ll *a,int N,int opt)
{
for(int i=1;i<N;i<<=1)
for(int p=i<<1,j=0;j<N;j+=p)
for(int k=0;k<i;++k)
/*or*/
if(opt==1)a[i+j+k]=(a[j+k]+a[i+j+k]);
else a[i+j+k]=(a[i+j+k]-a[j+k]);
/*and
if(opt==1)a[j+k]=(a[j+k]+a[i+j+k])%MOD;
else a[j+k]=(a[j+k]+MOD-a[i+j+k])%MOD;
*/
/*xor
{
int X=a[j+k],Y=a[i+j+k];
a[j+k]=(X+Y)%MOD;a[i+j+k]=(X+MOD-Y)%MOD;
//如果不是在模意义下的话，开一个long long，然后把逆元变成直接除二就好了
if(opt==-1)a[j+k]=1ll*a[j+k]*inv%MOD,a[i+j+k]=1ll*a[i+j+k]*inv%MOD;
}
*/
}
int getint()
{
int x=0,sign=1; char c=getchar();
while(c<'0' || c>'9')
{
if(c=='-') sign=-1; c=getchar();
}
while(c>='0' && c<='9')
{
x=x*10+c-'0'; c=getchar();
}
return x*sign;
}
void Init()
{
for(int i=0;i<=n;++i) Last[i]=0,Vis[i]=false;
tot=0;
}
void AddEdge(int u,int v)
{
Edge[++tot].to=v;
Edge[tot].last=Last[u]; Last[u]=tot;
}
void Read()
{
for(int i=1;i<=n;++i) Color[i]=getint()-1;
int u,v,w;
for(int i=1;i<n;i++)
{
u=getint(); v=getint();
AddEdge(u,v); AddEdge(v,u);
}
}
void GetRoot(int x,int father)
{
int v;
SonNum[x]=1; MaxNum[x]=1;
for(int i=Last[x];i;i=Edge[i].last)
{
v=Edge[i].to; if(v==father || Vis[v]) continue;
GetRoot(v,x);
SonNum[x]+=SonNum[v];
if(SonNum[v]>MaxNum[x]) MaxNum[x]=SonNum[x];
}
MaxNum[x]=max(MaxNum[x],ss-SonNum[x]);
if(rootx>MaxNum[x]) root=x,rootx=MaxNum[x];
}
int Deal(int x,int now)
{
return now|Power[Color[x]];
}
void GetDis(int x,int father,int now)
{
int t;
t=Record[++dlen]=Deal(x,now);
int v;
for(int i=Last[x];i;i=Edge[i].last)
{
v=Edge[i].to; if(v==father|| Vis[v]) continue;
GetDis(v,x,t);
}
}
ll Count(int x,int now)
{
for(int i=0;i<=dlen;++i) Record[i]=0;
dlen=0;
GetDis(x,0,now);
memset(Cnt,0,sizeof(Cnt));
for(int i=1;i<=dlen;++i)
++Cnt[Record[i]];
FWT(Cnt,len,1);
for(int i=0;i<len;++i) Cnt[i]=Cnt[i]*Cnt[i];
FWT(Cnt,len,-1);
ll ret=Cnt[len-1];
return ret;
}
void Solve(int x)
{
int v;
ans+=Count(x,0);
Vis[x]=true;
for(int i=Last[x];i;i=Edge[i].last)
{
v=Edge[i].to; if(Vis[v]) continue;
ans-=Count(v,Deal(x,0));
ss=SonNum[v]; rootx=INT_MAX; root=0;
GetRoot(v,x);
Solve(root);
}
}
void Work()
{
if(kk==1) return;
rootx=INT_MAX; ss=n; root=0; ans=0;
len=1<<kk;
GetRoot(1,0);
Solve(root);
}
void Write()
{
if(kk==1) printf("%lld\n",(ll)n*n);
else
printf("%lld\n",ans);
}
int main()
{
for(int i=0;i<15;i++) Power[i]=1<<i;
while(~scanf("%d%d",&n,&kk))
{
Init();
Read();
Work();
Write();
}
return 0;
}