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;
}