HDU-5977 Garden of Eden(树分治+枚举子集)

Garden of Eden

Time Limit: 10000/5000 MS (Java/Others)    Memory Limit: 131072/131072 K (Java/Others)
Total Submission(s): 673    Accepted Submission(s): 212

Problem Description

When God made the first man, he put him on a beautiful garden, the Garden of Eden. Here Adam lived with all animals. God gave Adam eternal life. But Adam was lonely in the garden, so God made Eve. When Adam was asleep one night, God took a rib from him and made Eve beside him. God said to them, “here in the Garden, you can do everything, but you cannot eat apples from the tree of knowledge.”
One day, Satan came to the garden. He changed into a snake and went to live in the tree of knowledge. When Eve came near the tree someday, the snake called her. He gave her an apple and persuaded her to eat it. Eve took a bite, and then she took the apple to Adam. And Adam ate it, too. Finally, they were driven out by God and began a hard journey of life.
The above is the story we are familiar with. But we imagine that Satan love knowledge more than doing bad things. In Garden of Eden, the tree of knowledge has n apples, and there are k varieties of apples on the tree. Satan wants to eat all kinds of apple to gets all kinds of knowledge.So he chooses a starting point in the tree,and starts walking along the edges of tree,and finally stops at a point in the tree(starting point and end point may be same).The same point can only be passed once.He wants to know how many different kinds of schemes he can choose to eat all kinds of apple. Two schemes are different when their starting points are different or ending points are different.

 

Input

There are several cases.Process till end of input.
For each case, the first line contains two integers n and k, denoting the number of apples on the tree and number of kinds of apple on the tree respectively.
The second line contains n integers meaning the type of the i-th apple. Types are represented by integers between 1 and k .
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

 

Output

For each case output your answer on a single line.

 

Sample Input

3 2 1 2 2 1 2 1 3

 

Sample Output

6

题意：树上有n个节点，每个节点有有k(k<=10)种苹果中的一个，定义一条合法路径为：路径上包含了k种苹果
问：合法路径的数量（起点或终点不相同则视为不同路径）

题解：树分治
首先状态压缩一下至(1< 而这道题处理的是状态，需要枚举状态进行合并（用莫比乌斯变换TLE了。。。）

设mx=(1< 我们也可以换个方向想：对于状态S，枚举S的子集x，再取SS=x^mx,这样只要合并S的子集x和SS
复杂度为O(3^n)

#include
using namespace std;
typedef long long LL;
const int MX = 5e4 + 5;
const int inf = 0x3f3f3f3f;
struct Edge {
    int v, nxt;
} E[MX * 2];
int tot, head[MX];
void init() {
    memset(head, -1, sizeof(head));
    tot = 0;
}
void add(int u, int v) {
    E[tot].v = v;
    E[tot].nxt = head[u];
    head[u] = tot++;
}

int n, k, a[MX];
int vis[MX], sz[MX], maxv[MX], root, Max, max_sta;
void dfs_size(int u, int fa) {
    sz[u] = 1; maxv[u] = 0;
    for (int i = head[u]; ~i; i = E[i].nxt) {
        int v = E[i].v;
        if (vis[v] || v == fa) continue;
        dfs_size(v, u);
        sz[u] += sz[v];
        maxv[u] = max(maxv[u], sz[v]);
    }
}
void dfs_root(int u, int fa, int rt) {
    maxv[u] = max(maxv[u], sz[rt] - maxv[u]);
    if (maxv[u] < Max) {
        Max = maxv[u];
        root = u;
    }
    for (int i = head[u]; ~i; i = E[i].nxt) {
        int v = E[i].v;
        if (vis[v] || v == fa) continue;
        dfs_root(v, u, rt);
    }
}
LL cnt_sta[1 << 10];
int Sta[15];
vectorv;
void dfs_sta(int u, int fa, int sta) {
    v.push_back(sta);
    for (int i = head[u]; ~i; i = E[i].nxt) {
        int v = E[i].v;
        if (vis[v] || v == fa) continue;
        dfs_sta(v, u, sta | Sta[a[v]]);
    }
}
LL cal(int rt, int sta) {
    v.clear();
    for (int i = 0; i <= max_sta; i++) cnt_sta[i] = 0;
    dfs_sta(rt, -1, sta | Sta[a[rt]]);
    for (int i = 0; i < v.size(); i++) cnt_sta[v[i]]++;
    LL ret = 0;
    for (int i = 0; i < v.size(); i++) {
        int S = v[i];
        for (int x = S; x; x = (x - 1)&S) ret += cnt_sta[x ^ max_sta];
        ret += cnt_sta[max_sta]; //x不能枚举到空集，空集需要特殊处理
    }
    return ret;
}

LL ans;
void DFS(int u) {
    Max = n;
    dfs_size(u, -1);
    dfs_root(u, -1, u);
    int rt = root;
    ans += cal(rt, 0);
    vis[rt] = 1;
    for (int i = head[rt]; ~i; i = E[i].nxt) {
        int v = E[i].v;
        if (vis[v]) continue;
        ans -= cal(v, Sta[a[rt]]);
        DFS(v);
    }
}

int main() {
    for (int i = 0; i < 11; i++) Sta[i] = 1 << i;
    //freopen("in.txt","r",stdin);
    while (~scanf("%d%d", &n, &k)) {
        init();
        max_sta = Sta[k] - 1;
        for (int i = 1; i <= n; i++) scanf("%d", &a[i]), a[i]--;
        for (int i = 1; i < n; i++) {
            int u, v;
            scanf("%d%d", &u, &v);
            add(u, v); add(v, u);
        }
        if (k == 1) {
            printf("%lld\n", (LL)n * n);
            continue;
        }
        memset(vis, 0, sizeof(vis));
        ans = 0;
        DFS(1);
        printf("%lld\n", ans);
    }
    return 0;
}