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 ci, and the i-th
12502
edge connects vertices ai and bi.
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(ai,bi) and C(bi,ai)
for all 1≤i≤(n-1). (i.e. |C(ai,bi)∩C(bi,ai)|)

Input

The input contains at most 15 sets. For each set:
The first line contains an integer n (2≤n≤105).
The second line contains n integers c1,c2,…,cn (1≤c_i≤n).
The i-th of the last (n-1) lines contains 2 integers ai,bi (1≤ai,bi≤n).

Output

For each set, (n-1) integers R1,R2,…,Rn-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就可以了。

        但是具体的，这样做相当于要维护到每个点的时候，以它为根的子树中包含颜色的种类和每种颜色的数量。然后所有儿子的这个颜色种类和数量要继承给父亲。这样就涉及到一个合并的问题。我们知道，普通的合并复杂度是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
,cnt
,ans
,n;
typedef pair<int,int> P;
vector<P> g
;
int root
;

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