hdu5977 Garden of Eden

都不好意思写题解了
跑了4000多ms
纪念下自己A的第二题
(我还有一道freetour II wa20多发没A。。。呜呜呜

#include<bits/stdc++.h>
using namespace std;
#define sz(X) ((int)X.size())
#define lson l,m,rt<<1
#define rson m+1,r,rt<<1|1
#define index Index
typedef long long ll;
const int N = 5e4+5;
const int INF = 0x3f3f3f3f;
const double pi = acos(-1.0);
int n,k,K;
ll ans;
int ty[N];
struct Node{
    int to,nx;
}E[N<<1];
int head[N], tot, vis[N];
void add(int u,int v) {
    E[tot].to = v; E[tot].nx = head[u]; head[u] = tot++;
}
/***************WeightRoot************/
int all, num, center;
int pp[N], nodes[N];
void findRoot(int x,int pre) {
    nodes[x] = 1; pp[x] = 0;
    for(int i = head[x]; ~i; i = E[i].nx) {
        int y = E[i].to; if(y == pre || vis[y]) continue;
        findRoot(y,x);
        nodes[x] += nodes[y];
        pp[x] = max(pp[x], nodes[y]);
    }
    pp[x] = max(pp[x], all-nodes[x]);
    if(pp[x] < num) {
        num = pp[x]; center = x;
    }
}
int getRoot(int root,int sn) {
    num = INF; all = sn; center = root;
    findRoot(root, -1);
    return center;
}
/****************treecdq**********/
ll has[1050];
ll dp[12][1050];
void getdp(int x, int pre, int num) {
    has[num] ++;
    for(int i = head[x]; ~i; i = E[i].nx) {
        int y = E[i].to; if(y == pre || vis[y]) continue;
        getdp(y,x, num|ty[y]);
    }
}
ll Cal(int x, int chu) {
    ll ret = 0;
    memset(has,0,sizeof(has));
    getdp(x,x,chu|ty[x]);
    for(int i = 0; i <= K; ++i) dp[0][i] = has[i];
    for(int i = 1; i <= k; ++i) {
        for(int j = 0; j <= K; ++j) {
            dp[i][j] = dp[i-1][j];
            if(!(j&(1<<(i-1)))) dp[i][j] += dp[i-1][j^(1<<(i-1))];
        }
    }
    for(int i = 0; i <= K; ++i) ret += has[i]* dp[k][i^K];
    return ret;
}
void work(int x) {
    vis[x] = 1;
    ans += Cal(x,0);
    for(int i = head[x]; ~i; i =  E[i].nx) {
        int y = E[i].to; if(vis[y]) continue;
        ans -= Cal(y,ty[x]);
        work(getRoot(y,nodes[y]));
    }
}
int main(){
    while(~scanf("%d %d",&n,&k)) {
        K = (1<<k)-1;
        memset(vis,0,sizeof(vis));
        memset(head,-1,sizeof(head)); tot = 0;
        for(int i = 1; i <= n; ++i) {
            int a; scanf("%d",&a); a--;
            ty[i] = 1<<a;
        }
        for(int i = 1; i < n; ++i) {
            int a,b; scanf("%d %d",&a,&b);
            add(a, b); add(b, a);
        }
        ans = 0;
        work(getRoot(1,n));
        printf("%lld\n", ans);
    }
    return 0;
}