CSU 1811 Tree Intersection(dfs序+莫队算法)

1811: Tree Intersection

Time Limit: 5 Sec   Memory Limit: 128 MB
Submit: 239   Solved: 108
[ Submit][ Status][ Web Board]

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

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

题意：一棵树，每个节点有一种颜色，现在删去一条边，求这个树分成两部分，这两部分的共有颜色个数。

解题思路：首先我们可以知道，一颗子树的颜色种类数为tot，而它所特有的颜色种类为c，则分开后的两部分树的共有颜色为tot-p。

由于是点权的更新，而且关于子树的部分，所以可以先把它的dfs序弄出来。并且这道题并不需要更新节点权值，所以可以用莫队算法，而且莫队算法处理这种题很好写代码。

#include<iostream>
#include<cstdio>
#include<cstring>
#include<vector>
#include<map>
#include<algorithm>
using namespace std;

const int maxn = 100005;
int n,sum[maxn],color[maxn],val[maxn];
int block,cnt,L[maxn],R[maxn],res[maxn];
vector<int> g[maxn];
int Map[maxn];
struct Edge
{
	int u,v;
}edge[maxn];
struct Query
{
	int l,r,id;

	bool operator < (const Query &rhs) const 
	{
		if(l / block == rhs.l / block)  
            return r < rhs.r;  
        return l / block < rhs.l / block; 
	}
}q[maxn];

void dfs(int u,int fa)
{
	L[u] = ++cnt;
	val[cnt] = color[u];
	for(int i = 0; i < g[u].size(); i++)
	{
		if(g[u][i] == fa) continue;
		dfs(g[u][i],u);
	}
	R[u] = cnt;
}

void solve()
{
	block = sqrt(cnt + 0.5);
	sort(q+1,q+n);
	int ans = 0,tot = 0,l = 1,r = 0;
	for(int i = 1; i < n; i++)
	{
		if(q[i].l == q[i].r)
		{
			if(sum[val[q[i].l]] == 1) res[q[i].id] = 0;
			else res[q[i].id] = 1;
			continue;
		}
		while(l < q[i].l)
		{
			Map[val[l]]--;
			if(Map[val[l]] == sum[val[l]] - 1) ans--;
			if(Map[val[l]] == 0) tot--;
			l++;
		}
		while(l > q[i].l)
		{
			l--;
			Map[val[l]]++;
			if(Map[val[l]] == sum[val[l]]) ans++;
			if(Map[val[l]] == 1) tot++;
		}
		while(r < q[i].r)
		{
			r++;
			Map[val[r]]++;
			if(Map[val[r]] == sum[val[r]]) ans++;
			if(Map[val[r]] == 1) tot++;
		}
		while(r > q[i].r)
		{
			Map[val[r]]--;
			if(Map[val[r]] == sum[val[r]] - 1) ans--;
			if(Map[val[r]] == 0) tot--;
			r--;
		}
		res[q[i].id] = tot - ans;
	}
	for(int i = 1; i < n; i++)
		printf("%d\n",res[i]);
}

int main()
{
	int u,v;
	while(scanf("%d",&n)!=EOF)
	{
		cnt = 0;
		memset(Map,0,sizeof(Map));
		memset(sum,0,sizeof(sum));
		for(int i = 1; i <= n; i++) g[i].clear();
		for(int i = 1; i <= n; i++)
		{
			scanf("%d",&color[i]);
			sum[color[i]]++;
		}
		for(int i = 1; i < n; i++)
		{
			scanf("%d %d",&u,&v);
			edge[i].u = u; edge[i].v = v;
			g[u].push_back(v);
			g[v].push_back(u);
		}
		dfs(1,-1);
		for(int i = 1; i < n; i++)
		{
			u = edge[i].u;
			v = edge[i].v;
			if(L[u] < L[v]) swap(u,v);
			q[i].l = L[u];
			q[i].r = R[u];
			q[i].id = i;
		}
		solve();
	}
	return 0;
}