CSU 1811 Tree Intersection（Treap启发式合并）

1811: Tree Intersection

Submit Page    Summary    Time Limit: 5 Sec     Memory Limit: 128 Mb     Submitted: 428     Solved: 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⁵).

The second line contains n integers c₁,c₂,…,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₁,R₂,…,R_n-1.

Sample Input

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

Sample Output

1211121

Hint

Source

湖南省第十二届大学生计算机程序设计竞赛

        非常气人的一道题目！对拍A了，CSUOJ却报我编译错误！

        人生中第一道启发式合并的题目……

        首先，这题题意是给你一棵树，然后每个节点都有一个颜色，然后问你如果从任意边开始截断，把树分为两个颜色的集合，问你交集的大小是多少。相当于就是计算每一条边的贡献。

        这题很类似于那道多校的题目，就是求路径上颜色数量，但是这题没有那么麻烦。我们考虑一个子树一个子树的处理。对于一个节点i，它与父亲之间的边的贡献，就是以i为根的子树中颜色数量小于整棵树数量的颜色种数。因为i点只有一个颜色，所以这个颜色种数就是它所有儿子的颜色种数之和，如果i点颜色加上去之后这种颜色数也等于总数则结果还要再加一。可以看出，我们只要dfs就可以了。

        但是具体的，这样做相当于要维护到每个点的时候，以它为根的子树中包含颜色的种类和每种颜色的数量。然后所有儿子的这个颜色种类和数量要继承给父亲。这样就涉及到一个合并的问题。我们知道，普通的合并复杂度是O(N^2)的，所以此时我们就需要用到数据结构来优化。这里，我使用了这两天一直在用的Treap。对于每个节点，我都建立一棵Treap，然后把每种颜色加入其中。处理完某一个点之后，再把该点Treap合并到其父亲那。合并的时候采取启发式合并，所谓启发式合并，就是每次合并的时候，把小的东西往大的东西上合并，这样可以保证插入的次数最小。可以证明，插入的次数不会超过logN。如此一来，合并的问题就得到了解决。

        具体来说的话维护Treap的时候有一些小技巧，比如说插入的时候给一个插入的位置，这样合并的时候就可以把原来废弃的点给用上，节省空间。然后处理个数的时候可以考虑用上取模，当某种颜色已经全部出现在子树中之后，它的数量对总数取模结果就是0，加上0对结果没有贡献。具体见代码吧（不要提交，玄学的CE啊~）：

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