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

4
1 2 2 1
1 2
2 3
3 4
5
1 1 2 1 2
1 3
2 3
3 5
4 5

Sample Output

1
2
1
1
1
2
1

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

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

线段树启发式合并：

#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 long
const 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 long
const 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;
}