CSU 1811 Tree Intersection（启发式合并）

思路：对于某一颗子树来说，只需要统计这颗子树拥有的颜色和这颗子树独有的颜色，两者一减就是答案，考虑一直DFS到子树，然后向上合并，合并的时候将结点数小的往大的合并，这样可以做到nlogn的复杂度

#include<bits/stdc++.h>using namespace std;const int maxn = 1e5+7;struct Edge{int v,id;Edge(){};Edge(int vv,int idd):v(vv),id(idd){};};vector<Edge>e[maxn];int n,c[maxn],sum[maxn],res[maxn],ans[maxn];map<int,int>cnt[maxn];    //cnt[u][i] 结点u颜色i有多少个void dfs(int u,int fa,int id){    cnt[u][c[u]]=1;ans[u]=cnt[u][c[u]]<sum[c[u]]?1:0;for(int i = 0;i<e[u].size();i++){int v = e[u][i].v;int id = e[u][i].id;if(v==fa)continue;dfs(v,u,id);        if(cnt[u].size()<cnt[v].size()){swap(ans[u],ans[v]);swap(cnt[u],cnt[v]);}for(map<int,int>::iterator it = cnt[v].begin();it!=cnt[v].end();it++){int &num = cnt[u][it->first];if(num==0 && num+it->second<sum[it->first])++ans[u];            else if (num && num+it->second==sum[it->first])--ans[u];num+=it->second;}}res[id]=ans[u];}int main(){while(scanf("%d",&n)!=EOF){memset(sum,0,sizeof(sum));for(int i = 0;i<=n;i++)e[i].clear(),cnt[i].clear();for(int i = 1;i<=n;i++)scanf("%d",&c[i]),sum[c[i]]++;for(int i = 1;i<=n-1;i++){int u,v;scanf("%d%d",&u,&v);e[u].push_back(Edge(v,i));e[v].push_back(Edge(u,i));}dfs(1,-1,0);        for(int i = 1;i<=n-1;i++)printf("%d\n",res[i]);} }

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 _iand 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