CSU 1811: Tree Intersection（线段树启发式合并||map启发式合并）

http://acm.csu.edu.cn/csuoj/problemset/problem?pid=1811

题意：给出一棵树，每一个结点有一个颜色，然后依次删除树边，问每次删除树边之后，分开的两个连通块里面的颜色交集数是多少，即公有的颜色数。

思路：可以像树形DP一样，先处理出儿子结点，然后回溯的时候和儿子结点的子树合并，更新父结点的子树，然后更新父结点的答案。

枚举删除的边。以u为子树里面放的是某个颜色有多少，然后和一开始统计的某种颜色的总数比较，如果树里面这个颜色的数目小于这个颜色的总数，那么这个颜色就肯定有一些是在另一个连通块里面，那么就是公有的，对答案有贡献。

如果树里面这个颜色的数目为0或者等于这个颜色的总数，说明不是公有的，那么对答案就没贡献。

因为直接合并的话O（n^2）的复杂度太大，因此用启发式合并达到O（nlogn）。

写了两种方法，都很好理解。

线段树：空间不足，因此要动态开辟结点

 #include <bits/stdc++.h>
 using namespace std;
 #define N 100010
 struct node {
     int val, cnt, l, r; // val是颜色c的个数,sum是答案的个数
 } tree[N*];
 struct Edge {
     int v, nxt, id;
 } edge[N*];
 int n, col[N], sum[N], head[N], tot, sz, root[N], ans[N], e[N];
 // 共有的就是交集就是答案
 void Add(int u, int v, int id) {
     edge[tot] = (Edge) { v, head[u], id }; head[u] = tot++;
     edge[tot] = (Edge) { u, head[v], id }; head[v] = tot++;
 }

 void PushUp(int now) {
     tree[now].cnt = tree[tree[now].l].cnt + tree[tree[now].r].cnt;
 }

 int Build(int l, int r, int c) {
     int now = ++sz;
     tree[now].l = tree[now].r = ;
     int m = (l + r) >> ;
     if(l == r) {
         tree[now].val = ;
         tree[now].cnt = (tree[now].val < sum[c] ?  : );
         return now;
     }
     if(c <= m) tree[now].l = Build(l, m, c);
     else tree[now].r = Build(m + , r, c);
     PushUp(now);
     return now;
 }

 void Merge(int &rt1, int rt2, int l, int r) {
     if(!rt1 || !rt2) {
         if(!rt1) rt1 = rt2; // 小的变成大的
         return ;
     }
     if(l == r) {
         tree[rt1].val += tree[rt2].val;
         tree[rt1].cnt = (tree[rt1].val < sum[l] ?  : );
         return ;
     }
     int m = (l + r) >> ;
     Merge(tree[rt1].l, tree[rt2].l, l, m);
     Merge(tree[rt1].r, tree[rt2].r, m + , r);
     PushUp(rt1);
 }

 void DFS(int u, int fa, int id) {
     root[u] = Build(, n, col[u]);
     for(int i = head[u]; ~i; i = edge[i].nxt) {
         int v = edge[i].v;
         if(v == fa) continue;
         DFS(v, u, edge[i].id);
         Merge(root[u], root[v], , n);
     }
     if(id) e[id] = tree[root[u]].cnt;
 }

 int main() {
     while(~scanf("%d", &n)) {
         memset(sum, , sizeof(sum));
         for(int i = ; i <= n; i++) scanf("%d", &col[i]), sum[col[i]]++;
         memset(head, -, sizeof(head)); sz = , tot = ;
         for(int i = ; i < n; i++) {
             int u, v; scanf("%d%d", &u, &v);
             Add(u, v, i);
         }
         DFS(, -, );
         for(int i = ; i < n; i++)
             printf("%d\n", e[i]);
     }
     return ;
 }

map：

 #include <bits/stdc++.h>
 using namespace std;
 #define N 100010
 struct Edge {
     int v, nxt, id;
 } edge[N*];
 map<int, int> num[N];
 int n, col[N], sum[N], cnt[N], e[N], head[N], tot;

 void Add(int u, int v, int id) {
     edge[tot] = (Edge) { v, head[u], id }; head[u] = tot++;
     edge[tot] = (Edge) { u, head[v], id }; head[v] = tot++;
 }

 void DFS(int u, int fa, int id) {
     num[u][col[u]] = ; cnt[u] = num[u][col[u]] < sum[col[u]] ?  : ;
     for(int i = head[u]; ~i; i = edge[i].nxt) {
         int v = edge[i].v, idd = edge[i].id;
         if(v == fa) continue;
         DFS(v, u, idd);
         if(num[u].size() < num[v].size()) // 启发式合并
             swap(num[u], num[v]), swap(cnt[u], cnt[v]);
         for(map<int,int>::iterator it = num[v].begin(); it != num[v].end(); it++) {
             int key = it->first, cc = it->second;
             if(num[u][key] + cc < sum[key] && num[u][key] == ) cnt[u]++; // 如果之前没被算过,并且是共有的就要加上
             if(num[u][key] + cc == sum[key] && num[u][key] > ) cnt[u]--; // 如果之前被算过,并且是特有的就要减去
             num[u][key] += cc;
         }
     }
     if(id) e[id] = cnt[u];
 }

 int main() {
     while(~scanf("%d", &n)) {
         memset(sum, , sizeof(sum));
         memset(cnt, , sizeof(cnt));
         for(int i = ; i <= n; i++) scanf("%d", &col[i]), sum[col[i]]++, num[i].clear();
         memset(head, -, sizeof(head)); tot = ;
         for(int i = ; i < n; i++) {
             int u, v; scanf("%d%d", &u, &v);
             Add(u, v, i);
         }
         DFS(, -, );
         for(int i = ; i < n; i++) printf("%d\n", e[i]);
     }
     return ;
 }