POJ-2114 : Boatherds

問題概要

頂点数V(<10^4)の重み付き木が与えられる。距離が丁度xに等しい2点が存在するか判定する問題。クエリはM(<100)個ある。

解法

前処理として重心分解して各重心ごとに部分木の距離を計算しておく。後は各クエリ毎に、部分木の距離配列からしゃくとりっぽくxを探せばよい。

acceptされたコード

重心分解はじめて書いてみたら詰まりに詰まったのでかなり蟻本を参考にした。あの実装頭よすぎる。

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

const int INF = int(1.05e9); 
const char *yes = "AYE", *no = "NAY";

typedef int cost_t;

struct Edge {
	int to;
	cost_t cost;

	Edge(){}
	Edge(int to, cost_t cost):
		to(to), cost(cost)
	{}
};


struct Graph : vector< vector<Edge> > {
	int numV;

	Graph(){}

	void init(int numV_) {
		numV = numV_;
		init();
	}

	void init() {
		assign(numV, vector<Edge>());
	}

	void addUndirectedEdge(int from, int to, cost_t cost) {
		operator [](from).push_back(Edge(to, cost));
		operator [](to).push_back(Edge(from, cost));
	}

	void addDirectedEdge(int from, int to, cost_t cost) {
		operator [](from).push_back(Edge(to, cost));
	}
};

const int MAX_Q = 100;
const int MAX_N = 10000;

int N, Q;
Graph G;
int qs[MAX_Q];
vector< pair<int,int> > dists[MAX_N];
bool isCenter[MAX_N];
int subTreeSize[MAX_N];

bool init() {
	scanf("%d", &N);
	if (N == 0) {
		return false;
	}
	memset(isCenter, false, sizeof(isCenter));
	for (int i = 0; i < N; ++i) {
		dists[i].clear();
	}
	G.init(N);
	for (int i = 0; i < N; ++i) {
		for (int x, c; scanf("%d", &x), x; scanf("%d", &c),G.addUndirectedEdge(i, x-1, c)) ;
	}
	Q = 0;
	for (int x; scanf("%d", &x), x; qs[Q++] = x) ;
	return true;
}

bool existX(const vector<pair<int,int> >& as, int x) {
	const int n = as.size();
	for (int i = 0, j = n - 1; i < j; ++i) {
		for (; j > i && as[i].first + as[j].first > x ; --j) ;
		if (j> i && as[i].second != as[j].second && as[i].first + as[j].first == x) {
			return true;
		}
	}
	return false;
}

int calcSize(int v, int par) {
	int ans = 1;
	for (int i = 0; i < int(G[v].size()); ++i) {
		const Edge &e = G[v][i];
		const int u = e.to;
		if (par != u && !isCenter[u]) {
			ans += calcSize(u, v);
		}
	}
	return subTreeSize[v] = ans;
}

pair<int,int> findCentroid(int v, int par, int sz) {
	pair<int,int> ret = make_pair(INF, -1);
	int subSize = 1, maxi = 0;
	for (int i = 0; i < int(G[v].size()); ++i) {
		const Edge &e = G[v][i];
		const int u = e.to;
		if (u != par && !isCenter[u]) {
			ret = min(ret, findCentroid(u, v, sz));
			maxi = max(maxi, subTreeSize[u]);
			subSize += subTreeSize[u];
		}
	}
	maxi = max(maxi, sz - subSize);
	ret = min(ret, make_pair(maxi, v));
	return ret;
}

void calcDistance(int v, int d, int type, int par, vector< pair<int,int> >& dist) {
	dist.push_back(make_pair(d, type));
	for (int i = 0; i < int(G[v].size()); ++i) {
		const Edge &e = G[v][i];
		const int u = e.to;
		if (u != par && !isCenter[u]) {
			calcDistance(u, d + e.cost, type, v, dist);
		}
	}
}

void decomposition(int v) {
	int sz = calcSize(v, -1);
	int center = findCentroid(v, -1, sz).second;
	isCenter[center] = true;

	dists[center].reserve(sz);
	dists[center].push_back(make_pair(0, -1));
	for (int i = 0; i < int(G[center].size()); ++i) {
		const Edge &e = G[center][i];
		const int u = e.to;
		if (!isCenter[u]) {
			calcDistance(e.to, e.cost, e.to, center, dists[center]);
		}
	}
	sort(dists[center].begin(), dists[center].end());

	for (int i = 0; i < int(G[center].size()); ++i) {
		const Edge &e = G[center][i];
		const int u = e.to;
		if (!isCenter[u]) {
			decomposition(u);
		}
	}
}

void solve() {
	decomposition(0);
	for (int i = 0; i < Q; ++i) {
		bool ans = false;
		for (int v = 0; v < N; ++v) {
			ans |= existX(dists[v], qs[i]);
		}
		puts(ans ? yes : no);
	}
	puts(".");
}

int main() {
	for (;init();) {
		solve();
	}
	return 0;
}