BUET Inter-University Programming Contest - 2011 A, UVa-12424 : Answering Queries on a Tree

問題概要

ノード数N(<10^5)の木があり、各頂点には色が塗られている(色の種類10)。M(<10^5)のクエリを処理する。クエリは2種類で、どれかの頂点の色を塗り替えるのと、2頂点間のパス上にもっとも多く現れる色の出現回数を求める。

解法

まず、適当な点を根にしてLCAを持ち、オイラーツアーを計算しておく。また、始めて訪ねたタイミングと出て行くときのタイミングを記録しておく。このとき、10種類のBITを用いて入るときに+1、出ていくときに-1を加えるようにすると根から任意の頂点までに何回出現するかを各色毎に計算できる。あとはLCAまでのを差っ引けばOK。

acceptされたコード

計算量O(MAX_COLOR * log N * M + N * log N)。

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

const int MAX_V = 100000;
const int MAX_C = 10;
const int INF = 1e9+10;

int V, Q;
int color[MAX_V];

struct RangeMinimumQuery{

	static const int MAX_RMQ = 2 * MAX_V - 1;
	static const int MAX_LOG = 20;

	int RMQ_SIZE;
	typedef int rmq_t;

	rmq_t rmq_data[MAX_RMQ];
	rmq_t sparse_table[MAX_LOG + 1][MAX_RMQ];

	void build_rmq(int N){
		RMQ_SIZE = N;
		const int log_rmq = 31 - __builtin_clz(RMQ_SIZE);

		for(int i=0; i<RMQ_SIZE; i++){
			sparse_table[0][i] = i;
		}

		for(int i=0; i<log_rmq; i++){
			for(int j=0, stp = RMQ_SIZE - (1<<(i+1)) + 1; j<stp; j++){
				const int prev = sparse_table[i][j], post = sparse_table[i][j+(1<<i)];
				sparse_table[i+1][j] = (rmq_data[prev] > rmq_data[post] ? post : prev);
			}
		}

	}

	//[from, to), indexを返す
	int query(int from, int to){
		const int log_d = 31 - __builtin_clz(to - from);
		const int prev = sparse_table[log_d][from], post = sparse_table[log_d][to-(1<<log_d)];
		return rmq_data[prev] > rmq_data[post] ? post : prev;
	}
};

RangeMinimumQuery rmq;

int vv;
int fs[MAX_V];
int ls[MAX_V];

int euler_tour[2*MAX_V - 1];
int depth[2*MAX_V - 1];
int first_occur[MAX_V];
vector<int> tree[MAX_V];

void dfs_LCA(int v, int par, int d, int& e){
	first_occur[v] = e;
	euler_tour[e] = v;
	depth[e++] = d;
	fs[v] = vv++;
	for(int i=0; i<(int)tree[v].size(); i++){
		const int u = tree[v][i];
		if(u != par){
			dfs_LCA(u, v, d + 1, e);
			euler_tour[e] = v;
			depth[e++] = d;
		}
	}
	ls[v] = vv++;
}

void build_LCA(int root){
	int e = 0;
	vv = 0;
	dfs_LCA(root, -1, 0, e);

	for(int i=0; i<2*V-1; i++){
		rmq.rmq_data[i] = depth[i];
	}
	rmq.build_rmq(2*V-1);
}

int query_LCA(int u, int v){
	return euler_tour[ rmq.query(min(first_occur[u], first_occur[v]), max(first_occur[u], first_occur[v])+1) ];
}

struct BinaryIndexedTree{

	typedef int bit_t;

	static const int MAX_BIT = 2*MAX_V;
	bit_t data[MAX_BIT+1];
	int SIZE;

	void init(int size){
		memset(data, 0, sizeof(data));
		SIZE = size;
	}

	bit_t sum(int n){
		bit_t ret = 0;
		for(;n;n-=n&-n){
			ret += data[n];
		}
		return ret;
	}

	bit_t sum(int from, int to){
		return sum(to)-sum(from);
	}

	void add(int n, bit_t x){
		for(n++;n<=SIZE;n+=n&-n){
			data[n]+=x;
		}
	}
};

BinaryIndexedTree bitree[MAX_C];

void init(){
	scanf("%d%d", &V, &Q);
	for(int i=0; i<V; i++){
		tree[i].clear();
	}

	for(int i=0; i<V; i++){
		scanf("%d", color+i);
		color[i]--;
	}

	for(int i=0; i<V-1; i++){
		int x, y;
		scanf("%d%d", &x, &y);
		x--; y--;
		tree[x].push_back(y);
		tree[y].push_back(x);
	}

	build_LCA(0);

	for(int i=0; i<MAX_C; i++){
		bitree[i].init(vv);
	}

	for(int i=0; i<V; i++){
		bitree[color[i]].add(fs[i], 1);
		bitree[color[i]].add(ls[i], -1);
	}
}

void query(){
	int cmd, a, b;
	scanf("%d%d%d", &cmd, &a, &b);
	a--; b--;
	if(cmd == 0){
		bitree[color[a]].add(fs[a], -1);
		bitree[color[a]].add(ls[a], 1);
		color[a] = b;
		bitree[color[a]].add(fs[a], 1);
		bitree[color[a]].add(ls[a], -1);
	}
	else {
		int ans = 0;
		int lca = query_LCA(a, b);
		for(int i=0; i<MAX_C; i++){
			int cur = bitree[i].sum(fs[a]+1) + bitree[i].sum(fs[b]+1) - 2 * bitree[i].sum(fs[lca] + 1);
			if(color[lca] == i){
				cur++;
			}
			//printf("a:%d, b:%d, col:%d, cur:%d\n", a, b, i, cur);
			//printf("f:%d s:%d l:%d\n", bitree[i].sum(fs[a]+1), bitree[i].sum(fs[b]+1), bitree[i].sum(fs[lca]+1));
			ans = max(ans, cur);
		}
		printf("%d\n", ans);
	}
}

int main(){
	int T;
	scanf("%d", &T);
	for(;T--;){
		init();
		for(int i=0; i<Q; i++){
			query();
		}
	}

	return 0;
}