POJ 2114 Boatherds 树分治

题意：

给出一颗边带权的\(n\)个节点的树，问是否存在最短距离为\(k\)的点对。

分析：

最开始做的题是询问最短距离小于等于\(k\)的点对。

我第一反应是下面的方法一。

方法一：
先求出小于等于\(k\)的点对 和 小于\(k\)(也就是小于等于\(k-1\))的点对，然后相减得到等于\(k\)的点对的个数。

方法二：
直接修改之前的统计点对的函数。

容易知道，第一种方法每组查询跑了两遍，所以运行时间是方法二的两倍。

方法一的代码：

#include <cstdio>
#include <cstring>
#include <algorithm>
#include <vector>
#include <map>
#define MP make_pair
using namespace std;

void read(int& x) {
	x = 0;
	char c = ' ';
	while(c < '0' || c > '9') c = getchar();
	while('0' <= c && c <= '9') {
		x = x * 10 + c - '0';
		c = getchar();
	}
}

typedef pair<int, int> PII;
const int maxn = 10000 + 10;
const int INF = 0x3f3f3f3f;

struct Edge
{
	int v, w, nxt;
	Edge() {}
	Edge(int v, int w, int nxt): v(v), w(w), nxt(nxt) {}
};

int n, k, ans;
vector<int> d, d2;
bool del[maxn];

int ecnt, head[maxn];
Edge edges[maxn * 2];

void AddEdge(int u, int v, int w) {
	edges[ecnt] = Edge(v, w, head[u]);
	head[u] = ecnt++;
}

int fa[maxn], sz[maxn];

void dfs(int u) {
	sz[u] = 1;
	for(int i = head[u]; ~i; i = edges[i].nxt) {
		int v = edges[i].v;
		if(v == fa[u] || del[v]) continue;
		fa[v] = u;
		dfs(v);
		sz[u] += sz[v];
	}
}

PII findCenter(int u, int t) {
	PII ans(INF, -1);
	int m = 0;
	for(int i = head[u]; ~i; i = edges[i].nxt) {
		int v = edges[i].v;
		if(v == fa[u] || del[v]) continue;
		ans = min(ans, findCenter(v, t));
		m = max(m, sz[v]);
	}
	m = max(m, t - sz[u]);
	ans = min(ans, MP(m, u));
	return ans;
}

void getDist(int u, int p, int dist, vector<int>& d) {
	d.push_back(dist);
	for(int i = head[u]; ~i; i = edges[i].nxt) {
		int v = edges[i].v, w = edges[i].w;
		if(v == p || del[v]) continue;
		getDist(v, u, dist + w, d);
	}
}

int cntPiars(vector<int>& d) {
	int ans = 0;
	sort(d.begin(), d.end());
	int j = d.size();
	for(int i = 0; i < d.size(); i++) {
		while(j && d[i] + d[j-1] > k) j--;
		ans += j - (j > i ? 1 : 0);
	}
	return ans;
}

void solve(int u) {
	fa[u] = 0;
	dfs(u);
	int s = findCenter(u, sz[u]).second;
	del[s] = true;
	
	for(int i = head[u]; ~i; i = edges[i].nxt) {
		int v = edges[i].v;
		if(del[v]) continue;
		solve(v);
	}

	d.clear();
	d.push_back(0);
	for(int i = head[s]; ~i; i = edges[i].nxt) {
		int v = edges[i].v, w = edges[i].w;
		if(del[v]) continue;
		d2.clear();
		getDist(v, s, w, d2);
		ans -= cntPiars(d2);
		d.insert(d.begin(), d2.begin(), d2.end());
	}

	ans += cntPiars(d);
	del[s] = false;
}

int main()
{
	while(scanf("%d", &n) == 1 && n) {
		ecnt = 0;
		memset(head, -1, sizeof(head));
		for(int u = 1; u <= n; u++) {
			int v, w;
			for(read(v); v; read(v)) {
				read(w);
				AddEdge(u, v, w);
				AddEdge(v, u, w);
			}
		}

		for(read(k); k; read(k)) {
			ans = 0;
			solve(1);
			int pre = ans;
			ans = 0;
			k--;
			solve(1);
			printf("%s\n", pre - ans > 0 ? "AYE" : "NAY");
		}
		printf(".\n");
	}

	return 0;
}

方法二的代码：

#include <cstdio>
#include <cstring>
#include <algorithm>
#include <vector>
#include <map>
#define MP make_pair
using namespace std;

void read(int& x) {
	x = 0;
	char c = ' ';
	while(c < '0' || c > '9') c = getchar();
	while('0' <= c && c <= '9') {
		x = x * 10 + c - '0';
		c = getchar();
	}
}

typedef pair<int, int> PII;
const int maxn = 10000 + 10;
const int INF = 0x3f3f3f3f;

struct Edge
{
	int v, w, nxt;
	Edge() {}
	Edge(int v, int w, int nxt): v(v), w(w), nxt(nxt) {}
};

int n, k, ans;
vector<int> d, d2;
bool del[maxn];

int ecnt, head[maxn];
Edge edges[maxn * 2];

void AddEdge(int u, int v, int w) {
	edges[ecnt] = Edge(v, w, head[u]);
	head[u] = ecnt++;
}

int fa[maxn], sz[maxn];

void dfs(int u) {
	sz[u] = 1;
	for(int i = head[u]; ~i; i = edges[i].nxt) {
		int v = edges[i].v;
		if(v == fa[u] || del[v]) continue;
		fa[v] = u;
		dfs(v);
		sz[u] += sz[v];
	}
}

PII findCenter(int u, int t) {
	PII ans(INF, -1);
	int m = 0;
	for(int i = head[u]; ~i; i = edges[i].nxt) {
		int v = edges[i].v;
		if(v == fa[u] || del[v]) continue;
		ans = min(ans, findCenter(v, t));
		m = max(m, sz[v]);
	}
	m = max(m, t - sz[u]);
	ans = min(ans, MP(m, u));
	return ans;
}

void getDist(int u, int p, int dist, vector<int>& d) {
	d.push_back(dist);
	for(int i = head[u]; ~i; i = edges[i].nxt) {
		int v = edges[i].v, w = edges[i].w;
		if(v == p || del[v]) continue;
		getDist(v, u, dist + w, d);
	}
}

int cntPiars(vector<int>& d) {
	int ans = 0;
	sort(d.begin(), d.end());
	int i = 0, j = d.size() - 1;
	while(i < j) {
		if(d[i] + d[j] < k) i++;
		else if(d[i] + d[j] > k) j--;
		else {
			if(d[i] == d[j]) {
				ans += (j - i + 1) * (j - i) / 2;
				break;
			} else {
				int l = i, r = j;
				while(d[i] == d[l]) i++;
				while(d[j] == d[r]) j--;
				ans += (i - l) * (r - j);
			}
		}
	}
	return ans;
}

void solve(int u) {
	fa[u] = 0;
	dfs(u);
	int s = findCenter(u, sz[u]).second;
	del[s] = true;
	
	for(int i = head[u]; ~i; i = edges[i].nxt) {
		int v = edges[i].v;
		if(del[v]) continue;
		solve(v);
	}

	d.clear();
	d.push_back(0);
	for(int i = head[s]; ~i; i = edges[i].nxt) {
		int v = edges[i].v, w = edges[i].w;
		if(del[v]) continue;
		d2.clear();
		getDist(v, s, w, d2);
		ans -= cntPiars(d2);
		d.insert(d.begin(), d2.begin(), d2.end());
	}

	ans += cntPiars(d);
	del[s] = false;
}

int main()
{
	while(scanf("%d", &n) == 1 && n) {
		ecnt = 0;
		memset(head, -1, sizeof(head));
		for(int u = 1; u <= n; u++) {
			int v, w;
			for(read(v); v; read(v)) {
				read(w);
				AddEdge(u, v, w);
				AddEdge(v, u, w);
			}
		}

		for(read(k); k; read(k)) {
			ans = 0;
			solve(1);
			printf("%s\n", ans ? "AYE" : "NAY");
		}
		printf(".\n");
	}

	return 0;
}