CSU1811-Tree Intersection

Tree Intersection

Submit Page     Summary    Time Limit: 1 Sec     Memory Limit: 128 Mb     Submitted: 498     Solved: 176    

---

Description

Bobo has a tree with n vertices numbered by1,2,…,n and (n-1) edges. The i-th vertex has color c i, and the i-th edgeconnects vertices a i and b i.Let C(x,y) denotes the set of colors insubtree rooted at vertex x deleting edge (x,y).Bobo would like to know R_i which is thesize of intersection of C(a i,b i) and C(b i,a i) for all 1≤i≤(n-1). (i.e. |C(ai,b i)∩C(b i,a i)|)

Input

The input contains at most 15 sets. Foreach set:

The first line contains an integer n(2≤n≤10 5).

The second line contains n integers c 1,c2,…,c n (1≤c_i≤n).

The i-th of the last (n-1) lines contains 2integers a i,b i (1≤a i,b i≤n).

Output

For each set, (n-1) integers R 1,R 2,…,R n-1.

Sample Input

41 2 2 11 22 33 451 1 2 1 21 32 33 54 5

Sample Output

1211121

题意：给你一棵树，对于每条边，把这条边删了，树分成了两个集合，求这两个集合中共同的颜色数量

解题思路：莫队或线段树启发式合并，关键都是看一个集合是否达到一种颜色的总量

线段树启发式合并：

#include <iostream>#include <cstdio>#include <cstring>#include <string>#include <algorithm>#include <map>#include <cmath>#include <set>#include <stack>#include <queue>#include <vector>#include <bitset>#include <functional>using namespace std;#define LL long longconst int INF = 0x3f3f3f3f;int n,x,y,a[100009],ans[100009],cnt[100009];int s[100009],nt[200009],e[200009],v[200009],res;int ss[100009],L[4000009],R[4000009],f[4000009],g[4000009],tot;void make(int &k,int l,int r,int p){    k=++tot,L[k]=R[k]=0;    if(l==r) {g[k]=1,f[k]=g[k]%cnt[p];return ;}    int mid=(l+r)>>1;    if(p<=mid) make(L[k],l,mid,p);    else make(R[k],mid+1,r,p);    f[k]=f[L[k]]+f[R[k]];}void Merge(int &now,int pre,int l,int r){    if(!now||!pre) {now=now^pre;return ;}    if(l==r) {f[now]=(bool)(g[now]=(g[now]+g[pre])%cnt[l]);return ;}    int mid=(l+r)>>1;    Merge(L[now],L[pre],l,mid);    Merge(R[now],R[pre],mid+1,r);    f[now]=f[L[now]]+f[R[now]];}void dfs(int k,int fa){    make(ss[k],1,n,a[k]);    for(int i=s[k];~i;i=nt[i])    {        if(e[i]==fa) continue;        dfs(e[i],k);        ans[v[i]]=f[ss[e[i]]];        Merge(ss[k],ss[e[i]],1,n);    }}int main(){while(~scanf("%d",&n))    {        memset(s,-1,sizeof s);        memset(cnt,0,sizeof cnt);        tot=L[0]=R[0]=res=0;        for(int i=1;i<=n;i++) scanf("%d",&a[i]),cnt[a[i]]++;        for(int i=1;i<n;i++)        {            scanf("%d%d",&x,&y);            nt[res]=s[x],s[x]=res,e[res]=y,v[res++]=i;            nt[res]=s[y],s[y]=res,e[res]=x,v[res++]=i;        }        dfs(1,0);        for(int i=1;i<n;i++) printf("%d\n",ans[i]);    }return 0;}

莫队：

#include <iostream>#include <cstdio>#include <cstring>#include <string>#include <algorithm>#include <map>#include <cmath>#include <set>#include <stack>#include <queue>#include <vector>#include <bitset>#include <functional>using namespace std;#define LL long longconst int INF = 0x3f3f3f3f;int s[100010],nt[200010],e[200010],flag[200010],id[200010];int n,u,v,res,N;int w[100010],cnt[100010],sum[100010];int x[100010],L[100010],R[100010];struct node{    int L,R,id;} q[100010];int ans[100010],ans1,ans2;void dfs(int k,int fa){    x[++res]=w[k];    L[k]=res;    for(int i=s[k]; ~i; i=nt[i])    {        if(e[i]==fa) continue;        flag[i]=1;        dfs(e[i],k);    }    R[k]=res;}bool cmp(node a,node b){    if(a.L/N==b.L/N) return a.R<b.R;    return a.L/N<b.L/N;}void check1(int k){    if(cnt[x[k]]-sum[x[k]]==1) ans1--;    if(sum[x[k]]==0) ans2--;}void check2(int k){    if(sum[x[k]]==cnt[x[k]]) ans1++;    if(sum[x[k]]==1) ans2++;}int main(){    while(~scanf("%d",&n))    {        memset(cnt,0,sizeof cnt);        memset(s,-1,sizeof s);        memset(sum,0,sizeof sum);        for(int i=1; i<=n; i++) scanf("%d",&w[i]),cnt[w[i]]++;        res=0;        for(int i=1; i<n; i++)        {            scanf("%d%d",&u,&v);            nt[res]=s[u],s[u]=res,e[res]=v,id[res]=i,flag[res++]=0;            nt[res]=s[v],s[v]=res,e[res]=u,id[res]=i,flag[res++]=0;        }        res=0;        dfs(1,-1);        res=0;        for(int i=1; i<=n; i++)        {            for(int j=s[i]; ~j; j=nt[j])            {                if(!flag[j]) continue;                q[res].L=L[e[j]];                q[res].R=R[e[j]];                q[res++].id=id[j];            }        }        N=sqrt(n);        sort(q,q+n-1,cmp);        ans1=0,ans2=0;        int ll=1,rr=0;        for(int i=0; i<n-1; i++)        {            while(ll<q[i].L) sum[x[ll]]--,check1(ll),ll++;            while(ll>q[i].L) ll--,sum[x[ll]]++,check2(ll);            while(rr>q[i].R) sum[x[rr]]--,check1(rr),rr--;            while(rr<q[i].R) rr++,sum[x[rr]]++,check2(rr);            ans[q[i].id]=ans2-ans1;        }        for(int i=1; i<n; i++) printf("%d\n",ans[i]);    }    return 0;}