HDU5977 Garden of Eden 【FMT】【树形DP】

题目大意：求有所有颜色的路径数。

题目分析：
参考codeforces997C，先利用基的FMT的性质在$O(2^k)$做FMT，再利用只还原一位的特点在$O(2^k)$还原，不知道为什么网上都要点分治。

代码：

#include<bits/stdc++.h>
#define R register
using namespace std;

const int maxn = 50200;

int n,k,a[maxn],fa[maxn],cnt[maxn],head[maxn];
struct edge{int to,nxt;}edges[maxn];
long long f[maxn],pt[maxn];

void read(){
    for(R int i=1;i<=n;i++) pt[i] = head[i] = 0;
    for(R int i=1;i<=n;i++) scanf("%d",&a[i]),a[i]--;
    for(R int i=1;i<n;i++){
    int u,v; scanf("%d%d",&u,&v);
    if(u > v) swap(u,v);
    fa[v] = u; edges[i] = (edge){v,head[u]}; head[u] = i;
    }
}

int dr;
void dfs(int now){
    long long tot = 0;long long dz = 0;
    for(R int i=head[now];i;i=edges[i].nxt){
    int to = edges[i].to;
    dfs(to); dz += tot*f[to]; tot += f[to]; 
    }
    dz *= f[now]; dz += f[now]*tot; f[now] = (tot+1)*f[now];
    pt[now] += dr*dz;
}

void work(){
    for(R int i=0;i<(1<<k);i++){
    for(R int j=1;j<=n;j++) f[j] = ((1<<a[j])&i)?1:0;
    dr = ((k-cnt[i])&1)?-1:1;
    dfs(1);
    }
    long long ans = 0;
    for(R int i=1;i<=n;i++) ans += pt[i];
    ans*=2; ans += n*(k==1);
    printf("%lld\n",ans);
}

int main(){
    for(R int i=1;i<=50000;i++) cnt[i] = cnt[i>>1]+(i&1);
    while(scanf("%d%d",&n,&k) == 2){ read(); work(); }
    return 0;
}