HDU-5977 Garden of Eden

题面

又是一堆废话，不挂了

题目链接https://vjudge.net/problem/HDU-5977

数据规模

Each of the following n-1 lines contains two integers u and v,meaning there is one edge between u and v.1≤n≤50000, 1≤k≤10

题意

给定一颗树，上面有k种苹果，每个结点都有一种苹果，问能吃到全部k种苹果的路线有多少种

题解

首先k只有10，我们完全可以枚举苹果的种类，而且可以状压表示吃到了哪几种苹果。我们再getdis中就可以用状压来表示吃到的苹果的状态了，getsta预处理结束后，状态之间按位或为((1<<k)-1)的都可以成为答案，所以我们就反过来想，枚举每一个getsta中得到的状态，然后枚举它的子集j，看(（(1<<k)-1) oplus j)是否存在，存在则加上答案即可。注意：

1. 枚举子集不会枚举到空集，所以先加上全集的答案。

2. 自己和自己产生答案也有可能，所以要先把自身-1，防止错误统计，统计结束后再加回来

还有在这个题中，边权变成了点权，我们在getsta中传入的参数就变成了点权，第一次加入答案的时候，u要考虑，要传入((1<<c[u]))，第二次减去答案，则传入((1 << c[u]) | (1 << c[v]))

代码

#include <bits/stdc++.h>
using namespace std;
const int N = 5e5 + 50;
vector<int> G[N];
const int inf = 1e9;
int sze[N];
int msze[N];
int root;
int vis[N];
int S;
int maxx;
void getroot(int u, int f) {
    sze[u] = 1;
    msze[u] = 0;
    for (int i = 0; i < G[u].size(); i++) {
        int v = G[u][i];
        if (vis[v] || v == f) continue;
        getroot(v, u);
        sze[u] += sze[v];
        msze[u] = max(msze[u], sze[v]);
    }
    msze[u] = max(msze[u], S - sze[u]);
    if (msze[u] < maxx) {
        maxx = msze[u];
        root = u;
    }
}
typedef long long ll;
int c[N];
ll a[5050];
int k;
vector<int> sta;
void getsta(int u, int f, int s) {
    a[s]++; sta.push_back(s);
    for (int i = 0; i < G[u].size(); i++) {
        int v = G[u][i];
        if (vis[v] || v == f) continue;
        getsta(v, u, s | (1 << c[v]));
    }
}

ll calc(int u, int s) {
    sta.clear();
    memset(a, 0, sizeof(a));
    getsta(u, 0, s);
    ll ans = 0;
    for (int i = 0; i < sta.size(); i++) {
        a[sta[i]]--;//注意1
        ans += a[(1 << k) - 1];//注意2
        for (int j = sta[i]; j; j = (j - 1) & sta[i]) {
            ans += a[((1 << k) - 1) ^ j];
        }
        a[sta[i]]++;//注意3
    }
    return ans;
}
ll ans = 0;
void dfs(int u) {
    vis[u] = 1;
    ans += calc(u, 1 << c[u]);//注意4
    for (int i = 0; i < G[u].size(); i++) {
        int v = G[u][i];
        if (vis[v]) continue;
        ans -= calc(v, (1 << c[v]) | (1 << c[u]));//注意5
        S = sze[v];
        root = 0;
        maxx = inf;
        getroot(v, u);
        dfs(root);
    }
}
int main() {
    int n;
    while (~scanf("%d%d", &n, &k)) {
        for (int i = 1; i <= n; i++) {
            G[i].clear();
            scanf("%d", &c[i]);
            c[i]--;
        }
        for (int i = 1; i < n; i++) {
            int u, v;
            scanf("%d%d", &u, &v);
            G[u].push_back(v);
            G[v].push_back(u);
        }
        if (k == 1) {//别忘了特判
            printf("%lld
", (ll)n * n);
            continue;
        }
        memset(vis, 0, sizeof(vis));//别忘清空vis!!
        ans = 0;
        S = n;
        root = 0;
        maxx = inf;
        getroot(1, 0);
        dfs(root);
        printf("%lld
", ans);
    }
    
    return 0;
}