AtCoder [ABC401F] Add One Edge 3 题解

回家路上20min胡出来的，好玩捏。

思路

考虑连接两个点 \(a,b\) 会导致直径发生什么变化。其中 \(a\) 在树 \(1\) ， \(b\) 在树 \(2\) 。

1. 新的直径没有经过 \((a,b)\)，这代表着直径只在其中一个树内，直径即为两个子树的直径中较长的一条。
2. 新的直径经过了 \((a,b)\)，此时贪心的想，可知直径长度为： \(a\) 点到树 \(1\) 中最远的距离 + \(b\) 点到树 \(2\) 中最远的距离 + \(1\)。其中最后的 \(1\) 为边 \((a,b)\) 。

那么我们首先求出两个树的直径。然后我们考虑如何求一个点到树内最远的距离。

根据树的直径的性质，一个点到树内最远的距离的点是树直径端点之一，反过来，两个端点到一个点的距离的 \(\max\) 即为一个点到树内最远的距离。跑 \(\operatorname{dfs}\) 即可。

我们依次枚举树 \(1\) 中的每个点，这时考虑如何将第二个 \(\operatorname{Sigma}\) 优化成 \(\log\) 级别的。很简单，将树 \(2\) 中每个点到树内最远距离进行排序，每次二分即可，详细过程见代码。

code

#include <bits/stdc++.h>

using namespace std;

#define int long long

int n[3], ans;
vector <int> sons[3][200005];
void add(int tid, int u, int v) {sons[tid][u].push_back(v); sons[tid][v].push_back(u);}
int jd[200005], lzl, lzr, rzl, rzr, zlen[3], zy[3][200005], q[200005];
void dfs(int tid, int xb) {
	int len = sons[tid][xb].size();
	for (int i = 0; i < len; i++) {
		int son = sons[tid][xb][i];
		if (jd[son] != -1) continue;
		jd[son] = jd[xb] + 1; dfs(tid, son);
	} 
}
int fard(int tid, int xb) {
	for (int i = 1; i <= n[tid]; i++) jd[i] = -1; jd[xb] = 0; dfs(tid, xb);
	int res = 1, Max = -1145141;
	for (int i = 1; i <= n[tid]; i++) if (jd[i] > Max) Max = jd[i], res = i;
	return res;
}
void ql(int tid, int xb) {
	for (int i = 1; i <= n[tid]; i++) jd[i] = -1; jd[xb] = 0; dfs(tid, xb);
	for (int i = 1; i <= n[tid]; i++) zy[tid][i] = max(zy[tid][i], jd[i]), zlen[tid] = max(zlen[tid], zy[tid][i]);
}
int B_S(int l, int r, int s) {
	if (l == r) return l;
	int mid = (l + r + 1) / 2;
	if (zy[1][mid] <= s) return B_S(mid, r, s);
	return B_S(l, mid - 1, s);
}

signed main() {
	cin >> n[0]; for (int i = 1; i < n[0]; i++) {int u, v; cin >> u >> v; add(0, u, v);}
	cin >> n[1]; for (int i = 1; i < n[1]; i++) {int u, v; cin >> u >> v; add(1, u, v);}
	lzl = fard(0, 1); lzr = fard(0, lzl); rzl = fard(1, 1); rzr = fard(1, rzl);
	ql(0, lzl); ql(0, lzr); ql(1, rzl); ql(1, rzr);
	sort(zy[1], zy[1] + n[1] + 1);//排序
	for (int i = 1; i <= n[1]; i++) q[i] = q[i - 1] + zy[1][i];//前缀和
	for (int i = 1; i <= n[0]; i++) {
		int res = B_S(0, n[1], max(zlen[0], zlen[1]) - zy[0][i] - 1);//res指直径为第一种情况的个数
		ans += (max(zlen[0], zlen[1]) * res + (zy[0][i] + 1) * (n[1] - res) + q[n[1]] - q[res]);
	}
	cout << ans << endl;
	return 0;
}