csu oj 1811: Tree Intersection (启发式合并)

题目链接：http://acm.csu.edu.cn/OnlineJudge/problem.php?id=1811

给你一棵树，每个节点有一个颜色。问删除一条边形成两棵子树，两棵子树有多少种颜色是有相同的。

启发式合并，小的合并到大的中。类似的题目有http://codeforces.com/contest/600/problem/E

//#pragma comment(linker, "/STACK:102400000, 102400000")
#include <algorithm>
#include <iostream>
#include <cstdlib>
#include <cstring>
#include <cstdio>
#include <vector>
#include <cmath>
#include <ctime>
#include <list>
#include <set>
#include <map>
using namespace std;
typedef long long LL;
typedef pair <int, int> P;
const int N = 1e5 + 5;
struct Edge {
    int next, to, index;
}edge[N << 1];
int color[N], head[N], tot;
int sum[N], ans[N], res[N]; //sum[color]:颜色color节点个数, ans[u]表示u点及字节点的答案, res[edge]表示边的答案
map <int, int> cnt[N]; //cnt[u][color] 表示u点子树color颜色有多少个节点

void init(int n) {
    for(int i = 1; i <= n; ++i) {
        head[i] = -1;
        sum[i] = 0;
        cnt[i].clear();
    }
    tot = 0;
}

inline void add_edge(int u, int v, int id) {
    edge[tot].next = head[u];
    edge[tot].to = v;
    edge[tot].index = id;
    head[u] = tot++;
}

void dfs(int u, int pre, int id) {
    cnt[u][color[u]] = 1;
    ans[u] = cnt[u][color[u]] < sum[color[u]] ? 1 : 0;
    for(int i = head[u]; ~i; i = edge[i].next) {
        int v = edge[i].to;
        if(v == pre)
            continue;
        dfs(v, u, edge[i].index);
        if(cnt[u].size() < cnt[v].size()) {
            swap(cnt[u], cnt[v]);
            swap(ans[u], ans[v]);
        }
        for(auto it : cnt[v]) {
            int &num = cnt[u][it.first];
            if(num == 0 && num + it.second < sum[it.first]) {
                ++ans[u];
            } else if(num + it.second == sum[it.first] && num) { //说明此子树的it.first颜色节点个数已满
                --ans[u];
            }
            num += it.second;
        }
    }
    res[id] = ans[u];
}

int main()
{
    int n, u, v;
    while(scanf("%d", &n) != EOF) {
        init(n);
        for(int i = 1; i <= n; ++i) {
            scanf("%d", color + i);
            ++sum[color[i]];
        }
        for(int i = 1; i < n; ++i) {
            scanf("%d %d", &u, &v);
            add_edge(u, v, i);
            add_edge(v, u, i);
        }
        dfs(1, -1, 0);
        for(int i = 1; i < n; ++i) {
            printf("%d
", res[i]);
        }
    }
    return 0;
}