CSU 1811 Tree Intersection （树形DP思想+动态开点/主席树思想+启发式合并）

题目链接：http://acm.csu.edu.cn:20080/csuoj/problemset/problem?pid=1811

题目大意：

给定一颗树，
树上每个节点都有颜色，问
对于每条边，其两端子树颜色集合交集的大小。

 题目分析：

这道题我是看题解才明白的，
算是启发式合并的一个入门练习吧。
首先想说这道题其实混合了不少思想在里面，
我大致觉得有树形DP，线段树/主席树，动态开点，启发式合并这些。
首先分析：不妨以1为根，假如一颗子树中某一颜色出现的个数
少于该颜色总数且不为零，那么这个颜色对该棵子树就有贡献。
于是我们要对每个节点讨论所有颜色，不妨对每个节点开线段树，
但是对于两颗树的合并，可以用动态开点的思想来节省空间，
并且两颗树可以共用某些节点。

复杂度的计算融入了启发式合并的思想，详细计算过程我不大会。。。只会感觉T_T/

#include<iostream>
#include<cstring>
#include<algorithm>
#include<cmath>
using namespace std;

#define debug puts("YES");
#define rep(x,y,z) for(int (x)=(y);(x)<(z);(x)++)
#define ll long long

#define lrt int l,int r,int rt
#define lson l,mid,rt<<1
#define rson mid+1,r,rt<<1|1
#define root l,r,rt
#define mst(a,b) memset((a),(b),sizeof(a))
#define pii pair<int,int>
#define fi first
#define se second
#define mk(x,y) make_pair(x,y)
const int mod=1e9+7;
const int maxn=1e5+5;
const int ub=1e6;
const double inf=1e-4;
ll powmod(ll x,ll y){ll t; for(t=1;y;y>>=1,x=x*x%mod) if(y&1) t=t*x%mod; return t;}
ll gcd(ll x,ll y){return y?gcd(y,x%y):x;}
/*
题目大意：
给定一颗树，
树上每个节点都有颜色，问
对于每条边，其两端子树颜色集合交集的大小。

题目分析：
这道题我是看题解才明白的，
算是启发式合并的一个入门练习吧。
首先想说这道题其实混合了不少思想在里面，
我大致觉得有树形DP，线段树/主席树，动态开点，启发式合并这些。
首先分析：不妨以1为根，假如一颗子树中某一颜色出现的个数
少于该颜色总数且不为零，那么这个颜色对该棵子树就有贡献。
于是我们要对每个节点讨论所有颜色，不妨对每个节点开线段树，
但是对于两颗树的合并，可以用动态开点的思想来节省空间，
并且两颗树可以共用某些节点。

复杂度的计算融入了启发式合并的思想，详细计算过程我不大会。。。只会感觉T_T/
*/
int n,x,y;
int a[maxn];
///前式链向星
struct node{
    int u,nxt,id;
}e[maxn<<1];
int head[maxn<<1],tot=0;
void init(){
    mst(head,-1),tot=0;
}
void add(int x,int y,int id){
    e[tot]=node{y,head[x],id};
    head[x]=tot++;
}
///动态开点线段树
struct seg{
    int v,cnt,l,r;
}tree[maxn<<5];
int rt[maxn],ans[maxn],sum[maxn];///
void pushup(int o){
    tree[o].cnt=tree[tree[o].l].cnt+tree[tree[o].r].cnt;
}
void build(int& o,int l,int r,int v){
    o=++tot;
    tree[o].l=tree[o].r=0;
    if(l==r){
        tree[o].v=1;
        tree[o].cnt=(tree[o].v<sum[l])?1:0;
        return ;
    }
    int mid=l+r>>1;
    if(v<=mid) build(tree[o].l,l,mid,v);
    else build(tree[o].r,mid+1,r,v);
    pushup(o);
}
void Merge(int& rt1,int rt2,int l,int r){
    if(!rt1||!rt2){
        if(!rt1) rt1=rt2;///合并
        return;
    }
    if(l==r){
        tree[rt1].v+=tree[rt2].v;
        tree[rt1].cnt=(tree[rt1].v<sum[l])?1:0;
        return ;
    }
    int mid=l+r>>1;
    Merge(tree[rt1].l,tree[rt2].l,l,mid);
    Merge(tree[rt1].r,tree[rt2].r,mid+1,r);
    pushup(rt1);
}
void dfs(int u,int fa,int id){
    build(rt[u],1,n,a[u]);///debug
    for(int i=head[u];~i;i=e[i].nxt){
        int v=e[i].u;
        if(v==fa) continue;
        dfs(v,u,e[i].id);
        Merge(rt[u],rt[v],1,n);
    }
    if(id) ans[id]=tree[rt[u]].cnt;
}

int main(){
    while(scanf("%d",&n)!=EOF){
        init(),mst(sum,0),mst(rt,0);
        rep(i,1,n+1) scanf("%d",&a[i]),sum[a[i]]++;
        rep(i,1,n){
            scanf("%d%d",&x,&y);
            add(x,y,i),add(y,x,i);
        }
        tot=0;dfs(1,-1,0);///开始支持线段树的构造
        rep(i,1,n) printf("%d\n",ans[i]);
    }
    return 0;
}