hdu5977 树分治

hdu5977
Garden of Eden

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

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

Source
2016ACM/ICPC亚洲区大连站-重现赛（感谢大连海事大学）

Recommend
wange2014
题意：给一棵树，树上有K（K<=10）种苹果，问从任意点出发，到任意点结束，把路径上的所有苹果取到，能够取全K种苹果的路径数（或则说点对数，1->3与3->1属于不同路径）
树分治的一道好题，和poj1741类似，是学习了 poj1741后才看懂的树分治和本道题，我详细写的poj1741
这道题我们类似的先计算从重心（root根）出发到达每个子树节点时，得到的苹果种类状态（状态压缩一下），然后统计每种状态的数量有多少。
我们记录了每个点的状态，用到一个和poj1741类似的方法来计算。

我们可以直接算出每个节点的状态来，即图中红色二进制数字表示的，然后可以统计出每种状态的数量。
我们遍历每个点，拿出该点的状态，寻找和该点状态或运算能够得到T（（1<

 for (int s0 = sta[i]; s0; s0 = (s0 - 1) & sta[i])///枚举i节点状态子集
        {
            res += Hash[((1 << k) - 1) ^ s0];
        }

枚举子集的一段神奇代码，举个例子：
枚举1010的子集
—— 运算——– 子集 -子集减一
1010 & 1010 = 1010 1001
1001 & 1010 = 1000 0111
0111 & 1010 = 0010 0001
0001 & 1010 = 0000 结束；

/**此代码借鉴了某大神的，我看懂了后又分析的*/
#include <stdio.h>
#include <iostream>
#include <algorithm>
#include <math.h>
#include <queue>
#include <set>
#include <string.h>
#include <vector>
using namespace std;
typedef long long ll;
const int MAXN = 5e5 + 10;

int n, k, Max, root;
ll ans;
vector <int> tree[MAXN];
vector <int> sta;
int sz[MAXN], maxv[MAXN], a[MAXN];
ll Hash[1200];
bool vis[MAXN];

void init()
{
    memset(vis, false, sizeof(vis));
    for (int i = 1; i <= n; i++) tree[i].clear();
}

void dfs_size(int u, int pre)
{
    sz[u] = 1;
    maxv[u] = 0;
    int cnt = tree[u].size();
    for (int i = 0; i < cnt; i++)
    {
        int v = tree[u][i];
        if (v == pre || vis[v]) continue;
        dfs_size(v, u);
        sz[u] += sz[v];
        maxv[u] = max(maxv[u], sz[v]);
    }
}

void dfs_root(int r, int u, int pre)
{
    maxv[u] = max(maxv[u], sz[r] - sz[u]);
    if (Max > maxv[u])
    {
        Max = maxv[u];
        root = u;
    }
    int cnt = tree[u].size();
    for (int i = 0; i < cnt; i++)
    {
        int v = tree[u][i];
        if (v == pre || vis[v]) continue;
        dfs_root(r, v, u);
    }
}
void dfs_sta(int u, int pre, int s)
{
    sta.push_back(s);
    int cnt = tree[u].size();
    for (int i = 0; i < cnt; i++)
    {
        int v = tree[u][i];
        if (v == pre || vis[v]) continue;
        dfs_sta(v, u, s | (1 << a[v]));
    }
}

ll cal(int u, int s)
{
    ll res = 0;
    sta.clear();
    dfs_sta(u, -1, s);
    memset(Hash, 0, sizeof(Hash));
    int cnt = sta.size();
   /**暴力去跑
    for(int i=0;i<cnt;i++)
    {
        for(int j=i+1;j<cnt;j++)
        {
            if(sta[i]|sta[j]==((1<<k)-1))
                res+=2;
        }
    }*/
    for (int i = 0; i < cnt; i++) Hash[sta[i]]++;
    for (int i = 0; i < cnt; i++)
    {
        Hash[sta[i]]--;///减去自身值
        res += Hash[(1 << k) - 1];///i点到状态为((1<<k)-1)肯定符合，所以直接加上
        for (int s0 = sta[i]; s0; s0 = (s0 - 1) & sta[i])///枚举i节点状态子集
        {
            res += Hash[((1 << k) - 1) ^ s0];
        }
        Hash[sta[i]]++;
    }
    return res;
}

void dfs(int u)
{
    Max = n;
    dfs_size(u, -1);
    dfs_root(u, u, -1);
    ans += cal(root, (1 << a[root]));
    vis[root] = true;
    int cnt = tree[root].size(), rt = root;
    for (int i = 0; i < cnt; i++)
    {
        int v = tree[rt][i];
        if (vis[v]) continue;
        ans -= cal(v, (1 << a[rt]) | (1 << a[v]));
        dfs(v);
    }
}

int main()
{
    while (scanf("%d%d", &n, &k) == 2)
    {
        init();
        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);
            tree[u].push_back(v);
            tree[v].push_back(u);
        }
        if (k == 1)
        {
            printf("%d\n", n * n);
            continue;
        }
        ans = 0;
        dfs(1);
        printf("%lld\n", ans);
    }
    return 0;
}