{题解}[jzoj5050]【GDOI2017模拟一试4.11】颜色树

传送门

Analysis

出题人特意卡了空间 禁用状压 DP
然而k还是给的很小 怎么说也是搜索允许的范围了啊！
从最简单的考虑
1. 当 $k = 1$ 时 则 $Ans = size^2$
2. 若 $k = 2$ 则 $Ans = (size^2) - 只有黑色(size^2) - 只有白色(size^2)$

可以看出 这是个容斥
然后 各位客官随意统计一下

Code

#include <cstdio>
#include <cstring>
#include <iostream>
#include <algorithm>
#define oo 2139062143
#define sqr(x) ((ll)(x)*(x))
#define abs(x) (((x)>=0)?(x):(-(x)))
#define max(x,y) (((x)>(y))?(x):(y))
#define min(x,y) (((x)<(y))?(x):(y))
#define fo(i,x,y) for (int i = (x);i <= (y);++ i)
#define fd(i,x,y) for (int i = (x);i >= (y);-- i)
#define fm(i,x) for (int i = las[x];i;i = e[i].nex)
using namespace std;
typedef double db;
typedef long long ll;
const int N = 50500,K = 11;
ll ans;
int n,k,size,tot;
int col[N];
struct node{
    int to,nex;
}e[N * 2];
int las[N];
bool vis[N];
int read()
{
    char ch=' ';int q=0,w=1;
    for(;(ch!='-')&&((ch<'0')||(ch>'9'));ch=getchar());
    if(ch=='-')w=-1,ch=getchar();
    for(;ch>='0' && ch<='9';ch=getchar())q=q*10+ch-48;
    int n=q*w;return n;
}
void link(int x,int y)
{
    e[++ tot].to = y,
    e[tot].nex = las[x],
    las[x] = tot;
}
void fill(int x,int s)
{
    vis[x] = true,++ size;
    fm(tmp,x)
        if (!vis[e[tmp].to] && (s & col[e[tmp].to]) == 0)
            fill(e[tmp].to,s);
}
void dfs(int lev,int f,int s)
//s 禁用颜色 
{
    if (lev > k)
    {
        memset(vis,0,sizeof vis);
        fo(i,1,n)
        if (!vis[i] && (col[i] & s) == 0)       
        {
            size = 0;
            fill(i,s);
            ans += (ll)sqr(size) * f;
        }
        return;
    }
    dfs(lev + 1,f,s);
    dfs(lev + 1,-f,s | (1 << (lev - 1)));
}
int main()
{
    freopen("colortree.in","r",stdin);
    freopen("colortree.out","w",stdout); 
    n = read(),k = read();
    fo(i,1,n)
    {
        int x = read();
        col[i] = 1 << (x - 1);
    }
    fo(i,1,n-1)
    {
        int x = read(),y = read();
        link(x,y),link(y,x);
    }
    dfs(1,1,0);
    printf("%lld\n",ans);
}