BZOJ3910. 火车

BZOJ3910. 火车
Q3.3.2.4. 火车

LCA 求距离
并查集让点不被重复标记
rt 如右图
\ a 到 b 的距离为
lca dep[a]+dep[b]-2*dep[lca]
/ \ 每次如果这个点被标记过, 则跳过
fa fb 反之, 暴力将这条链上的点标记
/ / 并查集优化:
a b 将 a,b,fa,fb,lca,... 的父节点设为 rt
每次标记时从 find(a) 到 find(b) 开始即可
若标记了, 则 find(x) != x

点击查看代码

#include <iostream>
#include <stdio.h>
#include <string.h>
#include <algorithm>
#include <utility>
#include <array>
#include <vector>

using namespace std;

const int N = 5e5 + 5, M = N << 1, logN = 25;

int n, m, rt;
int h[N], e[M], nxt[M], idx;
int f[N][logN], dep[N];
int father[N];

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

void add(int a, int b) {
	e[++ idx] = b, nxt[idx] = h[a], h[a] = idx;
}

void dfs(int u) {
	for(int i = 1; i < logN; i ++)
		if(f[u][i - 1]) f[u][i] = f[f[u][i - 1]][i - 1];
		else break;
	for(int i = h[u]; i; i = nxt[i]) {
		int v = e[i];
		if(v == f[u][0]) continue;
		f[v][0] = u, dep[v] = dep[u] + 1;
		dfs(v);
	}
}

int LCA(int x, int y) {
	if(dep[x] < dep[y]) swap(x, y);
	for(int i = logN - 1; i >= 0; i --)
		if(dep[f[x][i]] >= dep[y]) x = f[x][i];
	if(x == y) return x;
	for(int i = logN - 1; i >= 0; i --)
		if(f[x][i] != f[y][i]) x = f[x][i], y = f[y][i];
	return f[x][0];
}

int main() {
	scanf("%d%d%d", &n, &m, &rt);
	for(int i = 1, a, b; i < n; i ++)
		scanf("%d%d", &a, &b), add(a, b), add(b, a);
	for(int i = 1; i <= n; i ++) father[i] = i;
	dep[rt] = 1, dfs(rt);
	int last = rt, now;
	long long res = 0;
	for(int i = 1; i <= m; i ++) {
		scanf("%d", &now);
		if(find(now) != now) continue;
		int a = last, b = now;
		int lca = LCA(a, b);
		res += dep[a] + dep[b] - 2 * dep[lca];
		a = find(a), b = find(b);
		while(a != b) {
			if(dep[a] < dep[b]) swap(a, b);
			a = father[a] = find(f[a][0]);
		}
		if(find(lca) == lca) father[lca] = f[lca][0];
		last = now;
	}
	printf("%lld\n", res);
	return 0;
}