hihoCoder 1156 彩色的树

http://hihocoder.com/problemset/problem/1156

描述

给定一棵n个节点的树，节点编号为1, 2, …, n。树中有n - 1条边，任意两个节点间恰好有一条路径。这是一棵彩色的树，每个节点恰好可以染一种颜色。初始时，所有节点的颜色都为0。现在需要实现两种操作:

1. 改变节点x的颜色为y；

2. 询问整棵树被划分成了多少棵颜色相同的子树。即每棵子树内的节点颜色都相同，而相邻子树的颜色不同。

输入

第一行一个整数T，表示数据组数，以下是T组数据。

每组数据第一行是n，表示树的节点个数。接下来n - 1行每行两个数i和j，表示节点i和j间有一条边。接下来是一个数q，表示操作数。之后q行，每行表示以下两种操作之一：

1. 若为"1"，则询问划分的子树个数。

2. 若为"2 x y"，则将节点x的颜色改为y。

输出

每组数据的第一行为"Case #X:"，X为测试数据编号，从1开始。

接下来的每一行，对于每一个询问，输出一个整数，为划分成的子树个数。

数据范围

1 ≤ T ≤ 20

0 ≤ y ≤ 100000

小数据

1 ≤ n, q ≤ 5000

大数据

1 ≤ n, q ≤ 100000

        比赛时纯暴力过了小数据。后来看到岛娘说说：A map = =。想了想，知道是怎么回事了。把树转化为有根树，为每个节点开一个map，map[i][j]存储的是以i为根的子树中，颜色为j的孩子有多少个。维护这样的信息就可以解题了，因为整棵树的“颜色相同的子树的棵树”等于“与父亲颜色不同的节点数”+1，详见代码。

#include <bits/stdc++.h>
using namespace std;

#define ll long long

const int maxn=100010;
int tote;  
int head[maxn];  
int to[maxn<<1];  
int pre[maxn<<1]; 
int vis[maxn];
int p[maxn];
int color[maxn];
map<int,int> mp[maxn];


void addedge(int u,int v){  
    to[tote]=v;  
    pre[tote]=head[u];  
    head[u]=tote++;  
    to[tote]=u;  
    pre[tote]=head[v];  
    head[v]=tote++;  
}  

void dfs(int u){

	vis[u]=1;
	int cnt=0;
	for(int i=head[u];~i;i=pre[i]){
		int v=to[i];
		if(vis[v])continue;
		p[v]=u;
		dfs(v);
		cnt++;
	}
	mp[u][0]=cnt;
}

void init(){
	memset(head,-1,sizeof(head));
	memset(vis,0,sizeof(vis));
	memset(color,0,sizeof(color));
	tote=0;
}

int main(){
	int t;
	cin>>t;
	int cas=0;
	while(t--){
		init();
		cas++;
		
		int n;
		cin>>n;
		for(int i=1;i<=n;i++)mp[i].clear();
		
		for(int i=1;i<n;i++){
			int u,v;
			scanf("%d%d",&u,&v);
			addedge(u,v);
		}
		
		dfs(1);
		printf("Case #%d:\n",cas);
		int q;
		cin>>q;
		int ans=1;
		for(int i=1;i<=q;i++){
			int op;
			scanf("%d",&op);
			if(op==1){
				printf("%d\n",ans);
			}else{
				int x,y;
				scanf("%d%d",&x,&y);
				
				//不是根，维护当前节点与父亲关系 
				int px=p[x];
				if(px){
					if(--mp[px][color[x]]==0){
						mp[px].erase(color[x]);
					}
					
					if(color[x]==color[px]){
						if(color[px]!=y){
							ans++;
						}
					}
					if(color[x]!=color[px]){
						if(color[px]==y)ans--;
					}
				}
				
				//维护当前节点与孩子关系 
				if(mp[x].count(color[x])){
					ans+=mp[x][color[x]];
				}
				if(mp[x].count(y)){
					ans-=mp[x][y];
				}
				//最后变色 
				if(px)++mp[px][y];
				color[x]=y;
			}
			
		}
		
	}
	return 0;
}