Garden of Eden

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

还是统计树上 一笔画 的问题，难在 怎样快速的solve() 一颗子树的 所有的合法情况

 

我们保留所有的点，然后枚举他可能获得全部点的 数量。

#include<bits/stdc++.h>

using namespace std;
typedef long long LL;
typedef unsigned long long ULL;

#define rep(i,a,b) for(int i=a;i<b;++i)
#define per(i,a,b) for(int i=b-1;i>=a;--i)

const int N=50010;

struct Edge {
    int u,v,nt;
    Edge(int _u=0,int _v=0,int _nt=0)
    {
        u=_u,v=_v,nt=_nt;
    }
} edge[N*2];
int head[N],cnt;
void add_edge(int u,int v)
{
    edge[cnt]=Edge(u,v,head[u]);
    head[u]=cnt++;
}


bool vis[N];
int n,K,sz[N];
int max_num,rt,tot;//最大子树的大小，子树祖先，子树大小
void get_rt(int u,int f)
{
    sz[u]=1;
    int t=0;
    for(int i=head[u]; i!=-1; i=edge[i].nt) {
        Edge& e=edge[i];
        if(vis[e.v]||e.v==f)continue;
        get_rt(e.v,u);
        sz[u]+=sz[e.v];
        t=max(t,sz[e.v]);
    }
    t=max(t,tot-sz[u]);
    if(t<max_num) {
        max_num=t,rt=u;
    }
}


vector<int> vec;
LL rec[1200];
int col[N];

//记录子树中 从根向下的颜色数
void get_col(int u,int f,int s)
{
    vec.push_back(s);
    rec[s]++;
    for(int i=head[u]; i!=-1; i=edge[i].nt) {
        Edge& e=edge[i];
        if(vis[e.v]||e.v==f)continue;
        get_col(e.v,u,s|col[e.v]);
    }
}

LL solve(int u,int s)
{

    vec.clear();
    memset(rec,0,sizeof(rec));

    get_col(u,0,s);
    int sz=vec.size();
    LL res=0;
    int ttt=(1<<K)-1;

    /*
    我们拿到一个状态，想知道一共有多少个数字可以和当前的数字 或 之后，可以得到所有的点。 还是分类的思想，枚举子集进行计数
    */
    rep(i,0,sz) {
        rec[vec[i]]--;
        res+=rec[ttt];//空集的对应项,
        //printf("\n");
        
        //枚举每棵子树所有的点统计情况  nlog(n)*(2^K)	
        for(int sta=vec[i]; sta; sta=(sta-1)&vec[i]) { //枚举当前状态的子集，但是枚举不了空集
            //printf("vec:%d sta:%d\n",vec[i],sta);
            int t=sta^ttt;
            res+=rec[t];
        }
        rec[vec[i]]++;
    }
    //rep(i,0,sz)printf("i:%d %d\n",i,vec[i]);
    //printf("res:%d\n",res);
    return res;
}

LL res;
void work(int u) //u一定是祖先
{
    vis[u]=1;
    res+=solve(u,col[u]);
    for(int i=head[u]; i!=-1; i=edge[i].nt) {
        Edge& e=edge[i];
        if(vis[e.v])continue;
        res-=solve(e.v,col[u]|col[e.v]);

        rt=e.v,max_num=sz[e.v],tot=sz[e.v];
        get_rt(e.v,rt);
        work(rt);
    }
}
/*
3 2
1 2 2
1 2
1 3
2 1
1 1
1 2
1 1
1
*/
int main()
{
    while(scanf("%d %d",&n,&K)==2) {
        if(!n&&!K)break;

        //printf("%lld\n",1LL*n*n);

        cnt=0;
        fill(head,head+n+1,-1);
        fill(vis,vis+n+1,0);

        rep(i,1,n+1) {
            scanf("%d",&col[i]);
            col[i]--;
            col[i]=1<<col[i];
        }

        rep(i,0,n-1) {
            int u,v,w;
            scanf("%d %d",&u,&v);
            add_edge(u,v);
            add_edge(v,u);
        }
        
        if(K==1) {
            printf("%lld\n",1LL*n*n);
        } else {
            res=0;
            max_num=n,rt=1,tot=n;
            get_rt(1,0);
            //    printf("rt:%d\n",rt);
            work(rt);
            //rep(i,1,n+1)printf("i:%d deep:%d\n",i,deep[i]);
            printf("%lld\n",res);
        }
    }
    return 0;
}