gym-102082，2018-2019, ICPC, Asia Yokohama Regional Contest 2018 J Colorful Tree

题意

给定一个\(n\)个点的树，每个点有个初始颜色\(c_i\)，有\(m\)次询问，询问有两种：

• \(U~x_k~y_k\)，将第\(x_k\)个点的颜色改为\(y_k\)。
• \(Q~y_k\)，找到一个边数最少的子图，使得这个子图包括所有颜色为\(y_k\)的点，输出这个子图的边数。

分析

对每种颜色的点用一个set来存，set的排序规则按照dfs序，能发现每修改一个点的颜色对该颜色的答案的贡献只和其在set中前后的两个点\(l,r\)有关，set中增加一个点\(x\)对答案的贡献为\(dis(l,x)+dis(x,r)-dis(l,r)\)，删除一个点对答案的贡献为\(dis(l,r)-dis(l,x)-dis(x,r)\)，能发现这样得到的每个颜色的答案再加上set中第一个点\(start\)和最后一个点\(end\)之间的距离\(dis(start,end)\)刚好是正确答案的两倍。

对与每个询问答案的操作输出\((ans[i]+dis(start,end))/2\)。

Code

#include<algorithm>
#include<iostream>
#include<cstring>
#include<iomanip>
#include<sstream>
#include<cstdio>
#include<string>
#include<vector>
#include<bitset>
#include<queue>
#include<cmath>
#include<stack>
#include<set>
#include<map>
#define rep(i,x,n) for(int i=x;i<=n;i++)
#define per(i,n,x) for(int i=n;i>=x;i--)
#define sz(a) int(a.size())
#define rson mid+1,r,p<<1|1
#define pii pair<int,int>
#define lson l,mid,p<<1
#define ll long long
#define pb push_back
#define mp make_pair
#define se second
#define fi first
using namespace std;
const double eps=1e-8;
const int mod=1e9+7;
const int N=1e5+10;
const int inf=1e9;
int n,m;
vector<int>g[N];
int c[N],d[N],f[N],p[N],sz[N],top[N],son[N],ans[N],tot;
struct cmp{
	bool operator() (const int &x,const int &y)const{
		return p[x]<p[y];
	}
};
set<int,cmp>st[N];
void dfs(int u){
	d[u]=d[f[u]]+1;
	sz[u]=1;
	for(int x:g[u]){
		if(x==f[u]) continue;
		f[x]=u;
		dfs(x);
		sz[u]+=sz[x];
		if(sz[son[u]]<sz[x]) son[u]=x;
	}
}
void dfs1(int u,int t){
	p[u]=++tot;top[u]=t;
	if(son[u]) dfs1(son[u],t);
	for(int x:g[u]){
		if(x==f[u]||x==son[u]) continue;
		dfs1(x,x);
	}
}
int lca(int x,int y){
	while(top[x]!=top[y]){
		if(d[top[x]]<d[top[y]]) swap(x,y);
		x=f[top[x]];
	}
	if(d[x]<d[y]) swap(x,y);
	return y;
}
int dis(int x,int y){
	return d[x]+d[y]-2*d[lca(x,y)];
}
void add(int x){
	st[c[x]].insert(x);
	auto it=st[c[x]].find(x);
	int l=0,r=0;
	++it;
	if(it!=st[c[x]].end()){
		r=*it;
	}
	--it;
	if(it!=st[c[x]].begin()){
		--it;
		l=*it;
	}
	if(l&&r) ans[c[x]]-=dis(l,r);
	if(l) ans[c[x]]+=dis(l,x);
	if(r) ans[c[x]]+=dis(r,x);
}
void del(int x){
	auto it=st[c[x]].find(x);
	int l=0,r=0;
	++it;
	if(it!=st[c[x]].end()){
		r=*it;
	}
	--it;
	if(it!=st[c[x]].begin()){
		--it;
		l=*it;
	}
	if(l&&r) ans[c[x]]+=dis(l,r);
	if(l) ans[c[x]]-=dis(l,x);
	if(r) ans[c[x]]-=dis(r,x);	
	st[c[x]].erase(x);
}
int main(){
	ios::sync_with_stdio(false);
	//freopen("in","r",stdin);
	cin>>n;
	rep(i,2,n){
		int x,y;
		cin>>x>>y;
		g[x].pb(y);
		g[y].pb(x);
	}
	dfs(1);dfs1(1,1);
	rep(i,1,n){
		cin>>c[i];
		add(i);
	}
	cin>>m;
	while(m--){
		char cs;
		cin>>cs;
		int x,y;
		if(cs=='U'){
			cin>>x>>y;
			del(x);
			c[x]=y;
			add(x);
		}else{
			cin>>y;
			if(st[y].empty()) cout<<"-1\n";
			else cout<<(ans[y]+dis(*st[y].begin(),*st[y].rbegin()))/2<<'\n';
		}
	}
	return 0;
}