CSU 1811 Tree Intersection（Treap啟發式合并）

1811: Tree Intersection

​ Submit Page    Summary    5 Sec 128 Mb 428 161

Description

Bobo has a tree with n vertices numbered by 1,2,…,n and (n-1) edges. The i-th vertex has color c

i, and the i-th edge connects vertices a

i and b

i.

Let C(x,y) denotes the set of colors in subtree rooted at vertex x deleting edge (x,y).

Bobo would like to know R_i which is the size of intersection of C(a

i,b

i) and C(b

i,a

i) for all 1≤i≤(n-1). (i.e. |C(a

i,b

i)∩C(b

i,a

i)|)

Input

The input contains at most 15 sets. For each set:

The first line contains an integer n (2≤n≤10

5).

The second line contains n integers c

1,c

2,…,c

n (1≤c_i≤n).

The i-th of the last (n-1) lines contains 2 integers 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

Hint

Source

        非常氣人的一道題目！對拍A了，CSUOJ卻報我編譯錯誤！

        人生中第一道啟發式合并的題目……

        首先，這題題意是給你一棵樹，然後每個節點都有一個顔色，然後問你如果從任意邊開始截斷，把樹分為兩個顔色的集合，問你交集的大小是多少。相當于就是計算每一條邊的貢獻。

        這題很類似于那道多校的題目，就是求路徑上顔色數量，但是這題沒有那麼麻煩。我們考慮一個子樹一個子樹的處理。對于一個節點i，它與父親之間的邊的貢獻，就是以i為根的子樹中顔色數量小于整棵樹數量的顔色種數。因為i點隻有一個顔色，是以這個顔色種數就是它所有兒子的顔色種數之和，如果i點顔色加上去之後這種顔色數也等于總數則結果還要再加一。可以看出，我們隻要dfs就可以了。

#include<algorithm>
#include<iostream>
#include<cstring>
#include<cstdio>
#include<vector>
#include<cmath>
#define N 100010
#define file(x) freopen(#x".in","r",stdin),freopen(#x".out","w",stdout);
using namespace std;

int col[N],cnt[N],ans[N],n;
typedef pair<int,int> P;
vector<P> g[N];
 int root[N];

struct Treap
{
    struct treap
    {
        int son[2],col,cnt,fix,sum;
    } tree[N*10];

    inline void update(int x)
    {
        if (!x) return;
        tree[x].sum=bool(tree[x].cnt%=cnt[tree[x].col]);        //取模，如果用完則無貢獻，否則貢獻為1
        if (tree[x].son[0]) tree[x].sum+=tree[tree[x].son[0]].sum;
        if (tree[x].son[1]) tree[x].sum+=tree[tree[x].son[1]].sum;
    }

    inline void Rotate(int &x,bool ch)
    {
        int y=tree[x].son[ch^1],z=x;
        tree[x].son[ch^1]=tree[y].son[ch];
        tree[y].son[ch]=x; x=y;
        update(z); update(y);
    }

    inline void ins(int &i,int col,int t,int no)            //no表示插入的位置
    {
        if (!i)
        {
            i=no;
            tree[i].cnt=t;
            tree[i].col=col;
            tree[i].fix=rand();
            tree[i].sum=bool(t%cnt[col]);
            tree[i].son[0]=tree[i].son[1]=0;
            return;
        }
        if (col==tree[i].col) tree[i].cnt+=t;
        else
        {
            bool ch=(col>tree[i].col);
            ins(tree[i].son[ch],col,t,no);
            if (tree[tree[i].son[ch]].fix>tree[i].fix) Rotate(i,ch^1);
        }
        update(i);
    }

    inline void Merge(int &i,int j)                        //合并
    {
        if (!j) return;
        if (tree[j].son[0]) Merge(i,tree[j].son[0]);
        if (tree[j].son[1]) Merge(i,tree[j].son[1]);
        if (tree[j].cnt) ins(i,tree[j].col,tree[j].cnt,j);        //如果該點沒用完，還有意義，才需要合并。把廢棄的标号j利用上，省空間
    }
} treap;

void dfs(int x,int fa)
{
    treap.ins(root[x]=0,col[x],1,x);
    for(int i=0;i<g[x].size();i++)
    {
        int y=g[x][i].first;
        int id=g[x][i].second;
        if (y==fa) continue;
        dfs(y,x);
        ans[id]=treap.tree[root[y]].sum;
        if (ans[id]>treap.tree[root[x]].sum)            //啟發式合并，判斷顔色數量（已經全部出現過的顔色不用再計算）
        {
            treap.Merge(root[y],root[x]);
            root[x]=root[y];                        //别忘了換根
        }
        else treap.Merge(root[x],root[y]);
    }
}

void init()
{
    memset(g,0,sizeof(g));
    memset(cnt,0,sizeof(cnt));
}

int main()
{
    file(i);
    while(~scanf("%d",&n))
    {
        init();
        for(int i=1;i<=n;i++)
        {
            scanf("%d",&col[i]);
            cnt[col[i]]++;
        }
        for(int i=1;i<n;i++)
        {
            int u,v;
            scanf("%d%d",&u,&v);
            g[u].push_back(P(v,i));
            g[v].push_back(P(u,i));
        }
        dfs(1,0);
        for(int i=1;i<n;i++)
            printf("%d\n",ans[i]);
    }
    return 0;
}