「NOIP 2023 模拟赛 20230717 A」树上最长路 一种奇怪的做法

题意：
给你一棵 \(n\) 个节点的树, 每次询问包含第 \(i\) 条边的树上最长路径长度。
\(n\leq 10^6\)
赛时写了 2.5h，确实挺傻逼的。
首先，考虑先用 bfs 求出树的直径，然后用暴力 LCA 标记掉直径所有点。
然后对于所有直径上的节点 \(i\)，直接暴力 dfs 到所有非直径点，然后在这些子树上 DP 维护最长链 \(mxson_u\)，
即 \(mxson_u=\max\{mxson_v\}+1\)
同时维护每个点最近的直径点 \(F_u\) 和到直径点的距离 \(D_u\)。
然后对于每条边，假设离直径点更近的点是 \(a_i\)，答案为 \(mxson_{b_i}+D_{a_{i}}+\max(dist(s,F_{a_i}),dist(t,F_{a_i}))\)
比如说这颗树：

现在要求红边的答案。
那么可以分成绿色、蓝色、紫色三段：

绿色为 \(mxson_{b_i}\)，蓝色为 \(D_{a_i}\)，紫色为 \(\max(dist(s,F_{a_i}),dist(F_{a_i}))\)
那么 \(dist\) 在求直径的时候顺便求出来就好了。
总复杂度 \(\mathcal{O}(n)\)，但是跑得贼慢，吸氧加上 IO 优化才可以卡过 2s。

#include<bits/stdc++.h>
using namespace std;
#pragma GCC optimize(2)
#pragma GCC optimize(3)
#define I inline
#define gc getchar
const int N = 1000010; int h[N], dep[N], nxt[N * 2], to[N * 2], idx, n, a[N], b[N], d[N][3], s, t, fa[N]; bool vis[N], V[N]; int T, p[N], F[N], D[N], mxson[N];
I void add(int a, int b) {to[idx] = b; nxt[idx] = h[a]; h[a] = idx++;}
I void Read(int &x) {
	x = 0; char ch = gc(); while(ch < '0' || ch > '9') ch = gc();
	while(ch >= '0' && ch <= '9') x = x * 10 + (ch ^ 48), ch = gc();
}
I void Read(int &x, int &y) {Read(x), Read(y);}
I void write(int x) {if(x > 9) write(x / 10); putchar(x % 10 + '0');}
I void writeln(int x) {write(x); putchar('\n');}
I int bfs(int s, int f) {
	queue<int> q; for(int i = 1; i <= n; i++) vis[i] = 0; vis[s] = 1; d[s][f] = 0; q.push(s);
	while(!q.empty()) {
		int u = q.front(); q.pop();
		for(int i = h[u]; i != -1; i = nxt[i]) {
			int v = to[i]; if(vis[v]) continue; 
			vis[v] = 1; d[v][f] = d[u][f] + 1; q.push(v);}
		}
	
	int ans = s;
	for(int i = 1; i <= n; i++) if(d[i][f] > d[ans][f]) ans = i;
	return ans;
}
I void getfa(int u, int f) {p[u] = ++T; dep[u] = dep[f] + 1; fa[u] = f; for(int i = h[u]; i != -1; i = nxt[i]) {int v = to[i]; if(v == f) continue; getfa(v, u);}}
void tag(int a, int b) {
	if(dep[a] < dep[b]) swap(a, b); int delta = dep[a] - dep[b];
	while(delta--) vis[a] = 1, a = fa[a];
	if(a == b) return ;
	while(a != b) vis[a] = vis[b] = 1, a = fa[a], b = fa[b];
	vis[a] = 1;
}
int S;
I void dfs(int u, int f) {
	mxson[u] = 1; F[u] = S;
	for(int i = h[u]; i != -1; i = nxt[i]) {
		int v = to[i]; if(V[v] || v == f) continue;
		V[v] = 1; D[v] = D[u] + 1; dfs(v, u);
		mxson[u] = max(mxson[u], mxson[v] + 1);
	}
}
int main() {
//	freopen("tree.in", "r", stdin); freopen("tree.out", "w", stdout);
	Read(n); for(int i = 1; i <= n; i++) h[i] = -1; for(int i = 1; i < n; i++) {Read(a[i], b[i]); add(a[i], b[i]); add(b[i], a[i]);}
	s = bfs(1, 0); t = bfs(s, 1); int wyy = bfs(t, 2); 
	getfa(1, 0); for(int i = 1; i <= n; i++) vis[i] = 0; tag(s, t);
	for(int i = 1; i <= n; i++) if(vis[i]) V[i] = 1;
	for(int i = 1; i <= n; i++) if(vis[i]) S = i, dfs(i, 0);
	for(int i = 1; i < n; i++) {
		if(vis[a[i]] && vis[b[i]]) {
			writeln(d[t][1]);
		}
		else {
			if(D[a[i]] > D[b[i]]) swap(a[i], b[i]); // cout << "a = " << a[i] << " b = " << b[i] << '\n';
			writeln(mxson[b[i]] + D[a[i]] + max(d[F[a[i]]][1], d[F[a[i]]][2]));
		}
	}
	return 0;
}

虽然正解换根 DP 也是线性的，但是这种普及组做法好想，不过我写了 2.5h 实属傻逼
还好最后写出来了，不然彻底成消愁了。