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

 1 //#pragma comment(linker, "/STACK:102400000, 102400000")
 2 #include <algorithm>
 3 #include <iostream>
 4 #include <cstdlib>
 5 #include <cstring>
 6 #include <cstdio>
 7 #include <vector>
 8 #include <cmath>
 9 #include <ctime>
10 #include <list>
11 #include <set>
12 #include <map>
13 using namespace std;
14 typedef long long LL;
15 typedef pair <int, int> P;
16 const int N = 1e5 + 5;
17 struct Edge {
18     int next, to, index;
19 }edge[N << 1];
20 int color[N], head[N], tot;
21 int sum[N], ans[N], res[N]; //sum[color]:颜色color节点个数, ans[u]表示u点及字节点的答案, res[edge]表示边的答案
22 map <int, int> cnt[N]; //cnt[u][color] 表示u点子树color颜色有多少个节点
23 
24 void init(int n) {
25     for(int i = 1; i <= n; ++i) {
26         head[i] = -1;
27         sum[i] = 0;
28         cnt[i].clear();
29     }
30     tot = 0;
31 }
32 
33 inline void add_edge(int u, int v, int id) {
34     edge[tot].next = head[u];
35     edge[tot].to = v;
36     edge[tot].index = id;
37     head[u] = tot++;
38 }
39 
40 void dfs(int u, int pre, int id) {
41     cnt[u][color[u]] = 1;
42     ans[u] = cnt[u][color[u]] < sum[color[u]] ? 1 : 0;
43     for(int i = head[u]; ~i; i = edge[i].next) {
44         int v = edge[i].to;
45         if(v == pre)
46             continue;
47         dfs(v, u, edge[i].index);
48         if(cnt[u].size() < cnt[v].size()) {
49             swap(cnt[u], cnt[v]);
50             swap(ans[u], ans[v]);
51         }
52         for(auto it : cnt[v]) {
53             int &num = cnt[u][it.first];
54             if(num == 0 && num + it.second < sum[it.first]) {
55                 ++ans[u];
56             } else if(num + it.second == sum[it.first] && num) { //说明此子树的it.first颜色节点个数已满
57                 --ans[u];
58             }
59             num += it.second;
60         }
61     }
62     res[id] = ans[u];
63 }
64 
65 int main()
66 {
67     int n, u, v;
68     while(scanf("%d", &n) != EOF) {
69         init(n);
70         for(int i = 1; i <= n; ++i) {
71             scanf("%d", color + i);
72             ++sum[color[i]];
73         }
74         for(int i = 1; i < n; ++i) {
75             scanf("%d %d", &u, &v);
76             add_edge(u, v, i);
77             add_edge(v, u, i);
78         }
79         dfs(1, -1, 0);
80         for(int i = 1; i < n; ++i) {
81             printf("%d\n", res[i]);
82         }
83     }
84     return 0;
85 }