颜色树

题目大意

树上每个节点均有颜色。
求有多少条路径包含了所有k种颜色。k<=10

点分治

点分治后处理到根的二进制状态表示包含的颜色。
问题转化成插入若干个二进制数，查询包含某个二进制数的二进制数有多少个。
可以插入O(1)查询O(2^k)
考虑综合插入和查询的复杂度。
设f[A,B]表示有多少二进制数前一半为A后一半包含B。
插入时A是确定的，枚举包含的B，O(2^(k/2))
查询时B是确定的，枚举被包含的A，O(2^(k/2))

#include<cstdio>
#include<algorithm>
#define fo(i,a,b) for(i=a;i<=b;i++)
using namespace std;
typedef long long ll;
const int maxn=50000+10;
int h[maxn],go[maxn*2],next[maxn*2];
int d[maxn],size[maxn],a[maxn],v[maxn],f[40][40];
bool bz[maxn];
int i,j,k,l,t,n,m,tot,top,all;
ll ans;
int read(){
    int x=0,f=1;
    char ch=getchar();
    while (ch<'0'||ch>'9'){
        if (ch=='-') f=-1;
        ch=getchar();
    }
    while (ch>='0'&&ch<='9'){
        x=x*10+ch-'0';
        ch=getchar();
    }
    return x*f;
}
void add(int x,int y){
    go[++tot]=y;
    next[tot]=h[x];
    h[x]=tot;
}
void travel(int x,int y){
    a[++top]=x;
    int t=h[x];
    size[x]=1;
    while (t){
        if (!bz[go[t]]&&go[t]!=y){
            travel(go[t],x);
            size[x]+=size[go[t]];
        }
        t=next[t];
    }
}
void dfs(int x,int y){
    int t=h[x];
    while (t){
        if (!bz[go[t]]&&go[t]!=y){
            d[go[t]]=d[x]|v[go[t]];
            dfs(go[t],x);
        }
        t=next[t];
    }
}
void calc(int sig){
    int i,j,k,l,t;
    fo(i,1,top){
        t=d[a[i]];
        j=t/32;k=l=t%32;
        while (1){
            f[j][k]++;
            if (!k) break;
            k=(k-1)&l;
        }
    }
    fo(i,1,top){
        t=d[a[i]]^all;
        j=l=t/32;k=t%32;
        j^=31;l^=31;
        while (1){
            ans+=(ll)f[j^31][k]*sig;
            if (!j) break;
            j=(j-1)&l;
        }
    }
    fo(i,1,top){
        t=d[a[i]];
        j=t/32;k=l=t%32;
        while (1){
            f[j][k]--;
            if (!k) break;
            k=(k-1)&l;
        }
    }
}
void solve(int x,int y){
    top=0;
    travel(x,0);
    if (y) calc(-1);
    int i,j=x,k=0,t;
    while (1){
        t=h[j];
        while (t){
            if (!bz[go[t]]&&go[t]!=k&&size[go[t]]>top/2){
                k=j;
                j=go[t];
                break;
            }
            t=next[t];
        }
        if (!t) break;
    }
    d[j]=v[j];
    dfs(j,0);
    calc(1);
    bz[j]=1;
    t=h[j];
    while (t){
        if (!bz[go[t]]) solve(go[t],1);
        t=next[t];
    }
}
int main(){
    freopen("colortree.in","r",stdin);freopen("colortree.out","w",stdout);
    n=read();m=read();
    all=(1<<m)-1;
    fo(i,1,n) v[i]=read(),v[i]=1<<(v[i]-1);
    fo(i,1,n-1){
        j=read();k=read();
        add(j,k);add(k,j);
    }
    solve(1,0);
    printf("%lld\n",ans);
}