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 + ;
 struct Edge {
     int next, to, index;
 }edge[N << ];
 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 = ; i <= n; ++i) {
         head[i] = -;
         sum[i] = ;
         cnt[i].clear();
     }
     tot = ;
 }

 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]] = ;
     ans[u] = cnt[u][color[u]] < sum[color[u]] ?  : ;
     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 ==  && 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 = ; i <= n; ++i) {
             scanf("%d", color + i);
             ++sum[color[i]];
         }
         for(int i = ; i < n; ++i) {
             scanf("%d %d", &u, &v);
             add_edge(u, v, i);
             add_edge(v, u, i);
         }
         dfs(, -, );
         for(int i = ; i < n; ++i) {
             printf("%d\n", res[i]);
         }
     }
     return ;
 }