hdu 5977 Garden of Eden（点分治+状压dp）

#include <iostream>
#include <cstring>
#include <cstdio>
using namespace std;
typedef long long ll;
const int M = 5e4 + 10;
struct Edge {
    int v , next;
}edge[M << 1];
int head[M] , e , Size , root , n , k , a[M] , ssr;
ll ans;
bool vis[M];
void init() {
    memset(head , -1 , sizeof(head));
    memset(vis , false , sizeof(vis));
    ans = 0;
    e = 0;
    ssr = ((1 << k) - 1);
}
void add(int u , int v) {
    edge[e].v = v;
    edge[e].next = head[u];
    head[u] = e++;
}
int size[M] , mx[M];
ll Hash[1025];
void dfs_size(int u , int pre) {
    size[u] = 1;
    mx[u] = 0;
    for(int i = head[u] ; i != -1 ; i = edge[i].next) {
        int v = edge[i].v;
        if(v == pre || vis[v]) continue;
        dfs_size(v , u);
        size[u] += size[v];
        mx[u] = max(mx[u] , size[v]);
    }
}
void dfs_root(int r , int u , int pre) {
    mx[u] = max(mx[u] , size[r] - size[u]);
    if(mx[u] < Size) Size = mx[u] , root = u;
    for(int i = head[u] ; i != -1 ; i = edge[i].next) {
        int v = edge[i].v;
        if(v == pre || vis[v]) continue;
        dfs_root(r , v , u);
    }
}
void get_root(int u , int pre) {
    dfs_size(u , pre);
    dfs_root(u , u , pre);
}
int num , State[M];
void find_state(int u , int pre , int state) {
    State[num++] = state;
    for(int i = head[u] ; i != -1 ; i = edge[i].next) {
        int v = edge[i].v;
        if(vis[v] || v == pre) continue;
        find_state(v , u , state | (1 << a[v]));
    }
}
ll cau(int u , int state) {
    num = 0;
    find_state(u , -1 , state);
    memset(Hash , 0 , sizeof(Hash));
    ll sum = 0;
    for(int i = 0 ; i < num ; i++) Hash[State[i]]++;
    for(int i = 0 ; i < num ; i++) {
        Hash[State[i]]--;
        sum += Hash[ssr];//这里由于是枚举时0枚举不到所以先加上ssr^0.
        for (int s0 = State[i]; s0; s0 = (s0 - 1) & State[i])
            sum += Hash[((1 << k) - 1) ^ s0];
        Hash[State[i]]++;
    }
    return sum;
}
void dfs(int u) {
    Size = n;
    get_root(u , -1);
    ans += cau(root , (1 << a[root]));
    vis[root] = true;
    int rt = root;
    for(int i = head[root] ; ~i ; i = edge[i].next) {
        int v = edge[i].v;
        if(vis[v]) continue;
        ans -= cau(v , (1 << a[rt]) | (1 << a[v]));
        dfs(v);
    }
}
int main() {
    while(scanf("%d%d" , &n , &k) != EOF) {
        init();
        for(int i = 1 ; i <= n ; i++) scanf("%d" , &a[i]) , a[i]--;
        for(int i = 0 ; i < n - 1 ; i++) {
            int u , v;
            scanf("%d%d" , &u , &v);
            add(u , v);
            add(v , u);
        }
        if (k == 1) {
            printf("%lld\n", (ll)n * (ll)n);
            continue;
        }
        dfs(1);
        printf("%lld\n" , ans);
    }
    return 0;
}