AOJ 2170 - Marked Ancestor

問題 : Marked Ancestor

問題概要 :
大きさ$N$ の木が与えられる. 最初頂点$0$ は赤色で, その他の頂点は青色である. 次の$Q$ 個のクエリに答えよ :

・頂点$v$ の色を赤色に変える.
・頂点$v$ から赤色の先祖ノードのうち, もっとも近いものの番号をこたえる.

$1 \leqq N,\ Q \leqq 10^5$

解法 :

問題がXenia and Tree とほとんど被っている(これの簡単バージョン) なのでコードをほとんど流用できるが, 先祖ノードにしか巻き戻らないという事を活かしてもっとうまく解ける.

クエリを後ろから処理していくことを考えると, 色を変えるクエリは赤ｰ>青に戻すクエリと解釈できる. 色が戻っていくと, それぞれの頂点に対応するノードは部分木ごと併合されていくため, Union-Find 木を使うのが適していることが分かる.

ここで, 同じ頂点を何回も赤色を塗る事があるという点に気をつける.

コード :

#include <cstdio>
#include <cstring>
#include <vector>
#include <algorithm>

using namespace std;

vector<int> G[100000];
bool vis[100000];
int qtype[100000], vertex[100000], marked[100000];

int par[100000], pp[100000];

void init()
{
	for (int i = 0; i < 100000; i++) par[i] = i;
}

int find(int x)
{
	if (par[x] == x) return (x);
	return (par[x] = find(par[x]));
}

void merge(int x, int y)
{
	x = find(x); y = find(y);
	if (x != y){
		par[y] = x;
	}
}

void dfs(int v, int p, int m)
{
	int nm = m;
	if (marked[v]) nm = v;
	pp[v] = p;
	
	merge(nm, v);
	
	for (int i = 0; i < G[v].size(); i++){
		dfs(G[v][i], v, nm);
	}
}

int main()
{
	int N, Q;
	
	while (scanf("%d %d", &N, &Q) && N){
		for (int i = 0; i < N; i++) G[i].clear();
		
		for (int i = 1; i < N; i++){
			int v;
			scanf("%d", &v);
			G[--v].push_back(i);
		}
		
		init();
		
		memset(marked, 0, sizeof(marked));
		marked[0]++;
		for (int i = 0; i < Q; i++){
			char t[10];
			scanf("%s %d", t, vertex + i);
			qtype[i] = (t[0] == 'Q');
			--vertex[i];
			if (!qtype[i]) ++marked[vertex[i]];
		}
		
		dfs(0, -1, -1);
		
		long long int ans = 0;
		for (int i = Q - 1; i >= 0; i--){
			if (qtype[i]){
				ans += find(vertex[i]) + 1;
			}
			else {
				--marked[vertex[i]];
				if (!marked[vertex[i]]) merge(pp[vertex[i]], vertex[i]);
			}
		}
		
		printf("%lld\n", ans);
	}
	
	return (0);
}