ACM-ICPC 2018 Asia Yokohama Regional J Colorful Tree

問題

https://onlinejudge.u-aizu.ac.jp/resources/icpcooc2018/J.pdf

木が与えられる(頂点数100000まで)、各点には色が塗られている。
次の2つのクエリを最大100000回行う

  1. ある頂点の色を塗り替える。
  1. ある色に塗られている頂点が全て含まれるような最小の部分木の辺の数を出力する。

考察

各色について頂点集合を持つとすると、
1つ目のクエリについてはある色の頂点集合からその点を除き、他の色の頂点集合に加えるクエリになる。

1つの色に注目したとして、どうやって部分木を数えようかと色々書いているとLCAでできそうなのがわかる。

f:id:seica_at:20181215154925j:plain

上の写真において、例えば10と3に色が塗られていて、新しく11に色を塗るとするとこの色の答えは1増えることになるが、それを10と11のLCAと11との深さの差と見ることができる。

10と11に色が塗られていて、新しく13に色を塗ろうとすると答えが5増えることになるが、これは11と13のLCAと13との深さの差と、10と11のLCAと11と13のLCAとの深さの差の和と見ることができる。

上記のように考えると、木の番号をオイラーツアーで廻った順に振ると良いことがわかり。
各変更クエリについて頂点Aを付け足すときは

  • Aの番号の前後2つ頂点とのLCAと自身の深さの差のうち小さいほう
  • Aを付け足す前の頂点全体のLCAと付け加えた後の全体のLCAの深さの差

の2つを見ればよいことがわかる。

頂点を削除するときはこれと逆の操作をすればよい。計算量はO((n+m) log n)

#include <bits/stdc++.h>

using namespace std;

#define REP(i,n) for(ll (i) = (0);(i) < (n);++i)
#define PB push_back
#define MP make_pair
#define FI first
#define SE second
#define SHOW1d(v,n) {REP(WW,n)cerr << v[WW] << ' ';cerr << endl << endl;}
#define SHOW2d(v,WW,HH) {REP(W_,WW){REP(H_,HH)cerr << v[W_][H_] << ' ';cerr << endl;}cerr << endl;}
#define ALL(v) v.begin(),v.end()
#define Decimal fixed<<setprecision(20)
#define INF 1000000000
#define LLINF 1000000000000000000LL
#define MOD 1000000007

typedef long long ll;
typedef pair<ll,ll> P;

struct UF{
	vector<int> par;
	vector<int> sz;
	UF(int n):par(n),sz(n) {
		for(int i = 0; i < n; i++){
			par[i] = i;sz[i] = 1;
		}
	}
	int find(int x) {
		if (par[x] == x) return x;
		else return par[x] = find(par[x]);
	}
	void unite(int x, int y) {
		x = find(x); y = find(y);
		if (x == y) return;
		par[x] = y;
		sz[x] += sz[y];
	}
	bool same(int x, int y) { return find(x) == find(y); }
	int size(int n){return sz[find(n)];}
};

struct LCA{
	vector<vector<pair<int,int> > > g;
	vector<vector<int> > parent;
	vector<int> depth;
	UF uf;
	int sz;
	vector<bool> inited;
		
	LCA(int n):g(n),parent(20,vector<int>(n)),
	inited(n,false),depth(n),uf(n){sz = n;}
	
	void add_edge(int u,int v,int w = 1){
		if(uf.same(u,v))return;
		g[v].PB(MP(u,w));
		g[u].PB(MP(v,w));
		uf.unite(v,u);
	}
		
	void dfs(int v,int p,int d){
		inited[v] = true;
		parent[0][v] = p;
		depth[v] = d;
		REP(i,g[v].size()){
			if(g[v][i].FI != p)dfs(g[v][i].FI,v,d+1);
		}
	}
	
	void init(){
		REP(i,sz)if(!inited[i])dfs(uf.find(i),-1,0);
		REP(k,19){
			REP(i,sz){
				if(parent[k][i] < 0)parent[k+1][i] = -1;
				else parent[k+1][i] = parent[k][parent[k][i]];
			}
		}
	}
	
	int ret_lca(int a,int b){
		
		if(!uf.same(a,b))return -1;
		
		if(depth[a] > depth[b])swap(a,b);
		REP(k,20){
			if((depth[b] - depth[a])>>k&1)b = parent[k][b];
		}
		if(a == b)return a;
		
		for(int k = 19;k >= 0;k--){
			if(parent[k][a] != parent[k][b])
			{
				a = parent[k][a];
				b = parent[k][b];
			}
		}
		return parent[0][a];
	}
	
};

vector<vector<int>> v(111111);
int mp[111111];
int color[111111];
LCA lca(111111);

struct seica {
	set<pair<int,int>> s;
	int ans = 0;
	int top = 0;
	void add(int a){
		if(s.size() != 0){
			int ret = INF; 
			auto now = s.lower_bound(MP(a, lca.depth[a]));
			if(now != s.end()){
				int tmp = lca.ret_lca(a, now->FI);
				ret = min(ret, abs(lca.depth[a] - lca.depth[tmp]));
				tmp = lca.ret_lca(top, tmp);
				ans += abs(lca.depth[top] - lca.depth[tmp]);
				top = tmp;
			}
			if(now != s.begin()){
				now--;
				int tmp = lca.ret_lca(a, now->FI);
				ret = min(ret, abs(lca.depth[a] - lca.depth[tmp]));
				tmp = lca.ret_lca(top, tmp);
				ans += abs(lca.depth[top] - lca.depth[tmp]);
				top = tmp;
			}
			if(ret == INF)ret = 0;
			ans += ret;
		}
		else{
			top = a;
		}
		s.insert(MP(a, lca.depth[a]));
	}
	
	void remove(int a) {
		auto now = s.lower_bound(MP(a, lca.depth[a]));
		int ret = INF;
		now++;
		if(now != s.end()){
			int tmp = lca.ret_lca(a, now->FI);
			ret = min(ret, abs(lca.depth[a] - lca.depth[tmp]));
		}
		now--;
		if(now != s.begin()){
			now--;
			int tmp = lca.ret_lca(a, now->FI);
			ret = min(ret, abs(lca.depth[a] - lca.depth[tmp]));
			now++;
		}
		if(ret == INF)ret = 0;
		ans -= ret;
		s.erase(now);
		if(s.size() > 0){
			int a = s.begin() -> FI;
			auto hoge = s.end();hoge--;
			int b = hoge -> FI;
			int tmp = lca.ret_lca(a, b);
			ans -= abs(lca.depth[tmp] - lca.depth[top]);
			top = tmp;
		}
	}
};

seica ixmel[111111];
int sum;
void dfs(int a,int pre){
	REP(i,v[a].size()){
		if(v[a][i] == pre)continue;
		mp[v[a][i]] = sum;sum++;
		lca.add_edge(mp[a], mp[v[a][i]]);
		dfs(v[a][i], a);
	}
}

int main(){
	
	int n;cin >> n;
	REP(i,n-1){
		int a, b;cin >> a >> b;
		a--;b--;
		v[a].PB(b);
		v[b].PB(a);

	}
	
	REP(i,n)cin >> color[i];
	mp[0] = sum;sum++;
	dfs(0, -1);
	lca.init();
	REP(i,n)ixmel[color[i]].add(mp[i]);
	
	int q;cin >> q;
	REP(i,q){
		char c;cin >> c;
		if(c == 'U'){
			int a, b;cin >> a >> b;a--;
			ixmel[color[a]].remove(mp[a]);
			ixmel[b].add(mp[a]);
			color[a] = b;
		}
		else{
			int a;cin >> a;
			if(ixmel[a].s.size())cout << ixmel[a].ans << endl;
			else cout << -1 << endl;
		}
	}
	
	return 0;
}

最近コードに自分や人の名前を使いがち。