题目链接:http://acm.csu.edu.cn/csuoj/problemset/problem?pid=1811
题意:每个点有一个颜色,删除一条边后,求这条边两边的点集的颜色交的个数。
解法:平衡树启发式合并。那么可以直接用map对每个子树的信息进行记录,然后回溯上去的时候父节点加上子节点信息
的时候,把小的往大的加然后每次新出来一个颜色,就+1,如果颜色满了,就−1每次新搜一个节点的时候,如果子节点
更大,那么ans也要赋值成子节点的ans 想一下就知道了,知道启发式合并之后,难点就是ans的变化了
#include <bits/stdc++.h>
using namespace std;
const int maxn=100010;
struct edge{
int v,next,id;
}E[maxn*2];
int head[maxn],ans[maxn],sum[maxn],c[maxn],edgecnt;
map <int,int> cnt[maxn];
void init(){
memset(head,-1,sizeof(head));
edgecnt=0;
}
void add(int u,int v,int id){
E[edgecnt].v=v,E[edgecnt].next=head[u],E[edgecnt].id=id,head[u]=edgecnt++;
}
void dfs(int u, int fa, int id)
{
cnt[u][c[u]]=1;
if(sum[c[u]]>1) ans[id]++;
for(int i=head[u]; i+1; i=E[i].next){
int v=E[i].v;
if(v==fa) continue;
dfs(v,u,E[i].id);
if(cnt[u].size()<cnt[v].size()){
//swap(cnt[u],cnt[v]);
cnt[u].swap(cnt[v]);
ans[id]=ans[E[i].id];
}
for(map<int,int>::iterator it=cnt[v].begin(); it!=cnt[v].end(); it++){
int x=it->first,y=it->second;
if(cnt[u].count(x)){
cnt[u][x]+=y;
if(cnt[u][x]==sum[x]) ans[id]--;
}
else{
cnt[u][x]+=y;
if(cnt[u][x]<sum[x]) ans[id]++;
}
}
}
}
int main()
{
int n;
while(~scanf("%d", &n)){
memset(ans, 0, sizeof(ans));
memset(sum, 0, sizeof(sum));
for(int i=0; i<maxn; i++) cnt[i].clear();
for(int i=1; i<=n; i++) scanf("%d", &c[i]);
for(int i=1; i<=n; i++) sum[c[i]]++;
init();
for(int i=1; i<n; i++){
int u,v;
scanf("%d%d",&u,&v);
add(u,v,i);
add(v,u,i);
}
dfs(1,-1,0);
for(int i=1; i<n; i++) printf("%d\n", ans[i]);
}
return 0;
}