P8981 「DROI」Round 1 距离 题解
题意
给定一棵 \(n\) 个点的树，距离 \(\operatorname{dis}(u,v)\) 按点数 计。若点对 \((u,v)\) 满足任意点 \(x\) 到 \(u\)、\(v\) 的距离都不超过 \(\operatorname{dis}(u,v)\)，称之为极远点对 。设经过点 \(x\) 的极远点对数为 \(v_x\)，求
\(\sum_{x=1}^{n} v_x^k \bmod 998244353,\qquad k \in \{1,2\},\ n \leq 5 \times 10^6\)
显然：极远点对 = 树的直径
于是题目等价于：对每个点 \(x\)，统计有多少条直径经过 \(x\) 。注意星形树中有 \(\binom{n-1}{2} = \Theta(n^2)\) 条直径，所以只能计数，不能枚举。
Pre work
以下距离均用边数 ，记直径边数为 \(D\)。把树以 \(1\) 号点为根，需要三组信息，各由一次 DFS 完成：
- 直径长 \(D\)： 从 \(1\) 找最远点 \(a\)，再从 \(a\) 找最远点，最大距离即 \(D\)。
- 向下信息： \(\operatorname{dw}u\) 表示 \(u\) 到子树内最远点的边数，\(\operatorname{dc}u\) 为个数。后序遍历，用儿子的 \(\operatorname{dw}v+1\) 竞争最大值、并列累加个数即可。
- 向上信息： \(\operatorname{up}u\) 表示 \(u\) 到子树外最远点的边数，\(\operatorname{uc}u\) 为个数。自下而上的 DP 看不到子树外，必须换根：设 \(u\) 为 \(v\) 的父亲，则
\(\operatorname{up}v = 1 + \max\big(\operatorname{up}u,\ \max_{v' \neq v}(\operatorname{dw}v'+1),\ 0\big)\)
三个候选分别覆盖祖父方向、兄弟方向、父亲自身。实现时先在 \(u\) 处求出子方向链长的最大值 \((m_1,c_1)\) 与次大值 \((m_2,c_2)\)，再对每个儿子按"是否为最大值的唯一来源"决定用 \(m_1\) 还是 \(m_2\)。个数取并列达到最大值的候选之和。
核心：对每个点 \(x\)，统计经过它的直径数
把"经过 \(x\)"转成方向
从 \(x\) 出发的每条边通向一个方向：
- 儿子 \(v\) 方向：链长 \(\operatorname{dw}v + 1\)，个数 \(\operatorname{dc}v\)；
- 父亲方向：链长 \(\operatorname{up}x\)，个数 \(\operatorname{uc}x\)。
关键事实：一条路径经过 \(x\)，当且仅当它的两个端点落在 \(x\) 的两个不同方向 。若两端在同一方向，路径在子树内部就"拐弯"了，根本碰不到 \(x\)。
因此经过 \(x\) 的直径只有两种形态：
- 内点形态： 两端分属不同方向 \(i, j\)，且 \(d_i + d_j = D\)------两段链拼起来恰好是直径；
- 端点形态： \(x\) 自己是一端，某方向链长恰为 \(D\)，该方向每个最远点都与 \(x\) 构成一条直径。
设 \(x\) 的全部方向为 \((d_1, c_1), (d_2, c_2), \dots\)，则
\(\boxed{\,v_x = \underbrace{\sum_{i<j,\ d_i+d_j=D} c_i \cdot c_j}{\text{内点形态}} + \underbrace{\sum{i,\ d_i=D} c_i}_{\text{端点形态}}\,}\)
线性计算：边看边匹配
朴素枚举所有方向对是 \(O(\deg^2)\)，星形树的中心点会把总复杂度卡成 \(O(n^2)\)。发现：处理方向 \(i\) 时，真正需要的只是之前处理过的方向中，链长恰为 \(D - d_i\) 的个数之和，至于具体是哪些方向无关紧要。
维护按链长索引的桶 \(\operatorname{cntAt}d\)，逐方向处理：
-
先查询： \(v_x \mathrel{+} = c_i \cdot \operatorname{cntAt}D - d_i\)；
-
端点形态： 若 \(d_i = D\)，则 \(v_x \mathrel{+} = c_i\)；
-
后插入： \(\operatorname{cntAt}d_i \mathrel{+} = c_i\)。
-
注意：桶是全局数组，处理完 \(x\) 后按记录的已用下标逐个清零 ；若整体 memset，复杂度又退回 \(O(n^2)\)。
每个方向只做 \(O(1)\) 次操作，单点统计 \(O(\deg)\)，全树合计 \(O(\sum \deg) = O(n)\)
复杂度
\(O(n)\)
AC code
cpp
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
const ll MOD = 998244353;
const int N = 5e6 + 5;
ll n, k, ans;
int head[N], ver[N * 2], nxt[N * 2];
int tot;
ll dis[N];
ll dw[N], dc[N];
ll up[N], uc[N];
ll D;
void add(int x, int y) {
ver[++ tot] = y;
nxt[tot] = head[x];
head[x] = tot;
}
void pre(int u, int fa, int d) {
dis[u] = d;
for (int i = head[u]; i; i = nxt[i]) {
int v = ver[i];
if (v == fa) continue;
pre(v, u, d + 1);
}
}
void dfs1(int u, int fa) {
dw[u] = 0;
dc[u] = 1;
for (int i = head[u]; i; i = nxt[i]) {
int v = ver[i];
if (v == fa) continue;
dfs1(v, u);
ll val = dw[v] + 1;
ll cnt = dc[v];
if (val > dw[u]) {
dw[u] = val;
dc[u] = cnt;
}
else if (val == dw[u])
dc[u] = (dc[u] + cnt) % MOD;
}
}
ll cntAt[N], used[N];
ll usedTop;
void dfs2(int u, int fa) {
ll vx = 0;
usedTop = 0;
if (up[u] >= 0) {
if (D - up[u] >= 0) {
vx = (vx + (ll)uc[u] * cntAt[D - up[u]]) % MOD;
}
if (up[u] == D) {
vx = (vx + uc[u]) % MOD;
}
if (cntAt[up[u]] == 0) used[usedTop ++] = up[u];
cntAt[up[u]] = (cntAt[up[u]] + uc[u]) % MOD;
}
for (int i = head[u]; i; i = nxt[i]) {
int v = ver[i];
if (v == fa) continue;
ll val = dw[v] + 1;
ll cnt = dc[v];
if (D - val >= 0) {
vx = (vx + (ll)cnt * cntAt[D - val]) % MOD;
}
if (val == D) {
vx = (vx + cnt) % MOD;
}
if (cntAt[val] == 0) used[usedTop ++] = val;
cntAt[val] = (cntAt[val] + cnt) % MOD;
}
while (usedTop > 0) {
usedTop --;
cntAt[used[usedTop]] = 0;
}
if (k == 1) ans = (ans + vx) % MOD;
else ans = (ans + vx * vx % MOD) % MOD;
ll m1 = -1, c1 = 0, m2 = -1, c2 = 0;
for (int i = head[u]; i; i = nxt[i]) {
int v = ver[i];
if (v == fa) continue;
ll val = dw[v] + 1;
ll cnt = dc[v];
if (val > m1) {
m2 = m1;
c2 = c1;
m1 = val;
c1 = cnt;
}
else if (val == m1)
c1 = (c1 + cnt) % MOD;
else if (val > m2) {
m2 = val;
c2 = cnt;
}
else if (val == m2)
c2 = (c2 + cnt) % MOD;
}
for (int i = head[u]; i; i = nxt[i]) {
int v = ver[i];
if (v == fa) continue;
ll val = dw[v] + 1;
ll sibVal, sibCnt;
if (val != m1) {
sibVal = m1;
sibCnt = c1;
}
else if (c1 != dc[v]) {
sibVal = m1;
sibCnt = (c1 - dc[v] + MOD) % MOD;
}
else {
sibVal = m2;
sibCnt = c2;
}
ll bestVal = 0, bestCnt = 1;
if (up[u] >= 0) {
if (up[u] > bestVal) {
bestVal = up[u];
bestCnt = uc[u];
}
else if (up[u] == bestVal)
bestCnt = (bestCnt + uc[u]) % MOD;
}
if (sibVal >= 0) {
if (sibVal > bestVal) {
bestVal = sibVal;
bestCnt = sibCnt;
}
else if (sibVal == bestVal)
bestCnt = (bestCnt + sibCnt) % MOD;
}
up[v] = bestVal + 1;
uc[v] = bestCnt;
dfs2(v, u);
}
}
int main () {
ios::sync_with_stdio(false);
cin.tie(0);
cin >> n >> k;
tot = 0;
for (int i = 1; i < n; i ++) {
int u, v; cin >> u >> v;
add(u, v);
add(v, u);
}
pre(1, 0, 0);
int a = 1;
for (int i = 1; i <= n; i ++) {
if (dis[i] > dis[a]) a = i;
}
pre(a, 0, 0);
D = 0;
for (int i = 1; i <= n; i ++) D = max(D, dis[i]);
dfs1(1, 0);
up[1] = -1;
uc[1] = 0;
dfs2(1, 0);
cout << ans % MOD << '\n';
return 0;
}