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

//#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
, head
, tot;
int sum
, ans
, res
; //sum[color]:颜色color节点个数, ans[u]表示u点及字节点的答案, res[edge]表示边的答案
map <int, int> cnt
; //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\n", res[i]);
}
}
return 0;
}