CSU1811：Tree Intersection（树启发式合并）

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).

OutputFor each set, (n-1) integers R  ₁,R  ₂,…,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

题意：给一棵树，每个点有一种颜色，问对于每一条边，删除这条边之后形成的两棵子树的相交颜色种数。

思路：启发式合并，维护每个子树的颜色及其数量，分别讨论颜色的“由无到有”和“由有到满”两种情况就行了。这题还能用莫队算法，留坑。

# include <iostream>
# include <cstdio>
# include <map>
# include <vector>
# define PII pair<int,int>
using namespace std;
const int maxn = 1e5+30;
map<int,int>all, mp[maxn];
map<int,int>::iterator it;
int sum[maxn], ans[maxn], col[maxn];
vector<PII >g[maxn];
void dfs(int cur, int pre, int pid)
{
    for(int i=0; i<g[cur].size(); ++i)
    {
        int to = g[cur][i].second, id=g[cur][i].first;
        if(to == pre) continue;
        dfs(to, cur, id);
        if(mp[cur].size() < mp[to].size())
        {
            swap(mp[cur], mp[to]);
            swap(sum[pid], sum[id]);
        }
        for(it=mp[to].begin(); it!=mp[to].end(); ++it)
        {
            int color = it->first, num = it->second;
            if(!mp[cur].count(color))//无到有
            {
                if(num < all[color]) ++sum[pid];
                mp[cur][color] += num;
            }
            else//有到满
            {
                mp[cur][color] += num;
                if(mp[cur][color] == all[color]) --sum[pid];
            }
        }
    }
    if(!mp[cur].count(col[cur])) {if(++mp[cur][col[cur]] < all[col[cur]]) ++sum[pid];}
    else if(++mp[cur][col[cur]] == all[col[cur]]) --sum[pid];
    ans[pid] = sum[pid];
}
void init(int n)
{
    all.clear();
    for(int i=0; i<=n; ++i)
    {
        ans[i] = sum[i] = 0;
        mp[i].clear();
        g[i].clear();
    }
}
int main()
{
    int n, x;
    while(~scanf("%d",&n))
    {
        init(n);
        for(int i=1; i<=n; ++i)
        {
            scanf("%d",&col[i]);
            ++all[col[i]];
        }
        for(int i=1, u, v; i<n; ++i)
        {
            scanf("%d%d",&u,&v);
            g[u].push_back(make_pair(i,v));
            g[v].push_back(make_pair(i,u));
        }
        dfs(1, 0, 0);
        for(int i=1; i<n; ++i) printf("%d\n",ans[i]);
    }
    return 0;
}

<iframe src="http://free.timeanddate.com/clock/i67ujt5q/n459/szw110/szh110/hoc222/hbw6/cf100/hgr0/hcw2/hcd88/fan2/fas20/fdi70/mqc000/mqs3/mql13/mqw4/mqd94/mhc000/mhs3/mhl13/mhw4/mhd94/mmc000/mml5/mmw1/mmd94/hwm2/hhs2/hhb18/hms2/hml80/hmb18/hmr7/hscf09/hss1/hsl90/hsr5" frameborder="0" width="110" height="110"></iframe>