2.27省选模拟赛补题记录：直径(容斥，树形dp，换根dp)
题意
定义一棵树的直径条数为
(
n
2
)
\binom{n}{2}
(2n) 对点中，取道距离最大值的选法数量。
给定一棵
n
n
n 个点的树，你可以将每条边的权值赋值为
0
0
0 或
1
1
1。
你需要求出所有
2
n
−
1
2^{n - 1}
2n−1 种赋值方法生成的树的直径条数之和。
你只需要输出答案对
998244353
998244353
998244353 取模后的结果即可。
2 ≤ n ≤ 2000 2 \leq n \leq 2000 2≤n≤2000。
分析
很好的题，教会了我一些 t r i c k trick trick。
看到边权只有
0
/
1
0/1
0/1，考虑在 中点 对每条直径计数。
也就是说这
2
n
−
1
2^{n - 1}
2n−1 种赋值方案下，每条直径都会对应一个或多个中点(存在多个中点的原因是因为有边权为
0
0
0 的边)，考虑枚举中点计算所有方案下以它作为中点的直径条数之和。
那么中点可能在 点 上，也可能在 边 上。如果在边上就是边的中点位置。
考虑这两种情况下怎么算直径条数：
-
在点上
那么这个点是某条直径中点的充要条件：是以这个点为根时存在 ≥ 2 \geq 2 ≥2 儿子子树中含有最大深度的点。这里 x x x 的深度定义为 x x x 到根路径上的边权和。证明需要考虑一下满足上述条件时为什么会不存在一条不经过这个点但是比 2 × m a x d e p 2 \times maxdep 2×maxdep 还长的路径。 -
在边上
实际上跟上面是一样的，把这条边的两个端点看作两个根，那么这两个根的子树中的最大深度相同。
因此可以 d p dp dp，只需要记录子树内最大深度为 i i i 时的赋值方案数 f x , i f_{x, i} fx,i，以及在这些方案中能取到最大深度的叶子数量和 g x , i g_{x, i} gx,i。然后在根处合并一下。
但是现在有一个问题：由于边权可以赋值为
0
0
0，因此一种赋值方案下的一条直径可能被计算多次。
比如下面的例子：
这条直径会在枚举到
3
,
4
,
5
3,4,5
3,4,5 和
(
3
,
4
)
,
(
4
,
5
)
(3, 4), (4, 5)
(3,4),(4,5) 时被算到。
注意到一条直径被算到的位置一定由边权为
0
0
0 的边连成一个连通块。考虑 点 - 边 = 1 容斥。
那么枚举到点时就加上求出的直径条数。枚举到边时，如果这条边的边权为
0
0
0，就减去直径条数，如果边权为
1
1
1，就加上直径条数。
那么无论中点在某条
1
1
1 权边上，还是在若干点和
0
0
0 权边上，贡献都只会被算一次。
注意到对边的计算正负抵消，因此只用计算点的贡献就行了。
但是枚举根，一次
d
p
dp
dp 的复杂度是
O
(
n
2
)
O(n^2)
O(n2) 的，总复杂度就是
O
(
n
3
)
O(n^3)
O(n3)。
可以换根优化到
O
(
n
2
)
O(n^2)
O(n2)。我写的时空复杂度都是
O
(
n
2
)
O(n^2)
O(n2) 的。
CODE：
#include<bits/stdc++.h>
#define pb emplace_back
using namespace std;
typedef long long LL;
const int N = 2010;
const LL mod = 998244353;
int n, dep[N];
LL ans;
vector< int > E[N];
vector< LL > f[N], g[N];
struct node {
int x; vector< LL > v;
};
vector< node > F[N], G[N];
void dfs(int x, int fa) {
int c = 0;
for(auto v : E[x]) {
if(v == fa) continue;
c ++;
dfs(v, x); dep[x] = max(dep[x], dep[v] + 1);
F[x].pb((node) {v, f[v]});
G[x].pb((node) {v, g[v]});
f[v].pb(0); g[v].pb(0); // 补充最大深度
for(int i = f[v].size() - 1; i > 0; i -- ) {
f[v][i] = (f[v][i] + f[v][i - 1]) % mod;
g[v][i] = (g[v][i] + g[v][i - 1]) % mod;
}
}
if(!c) {f[x].pb(1); g[x].pb(1); return ;}
else {
for(auto v : E[x]) {
if(v == fa) continue;
for(int i = 1; i < f[v].size(); i ++ ) f[v][i] = (f[v][i] + f[v][i - 1]) % mod;
}
for(int i = 0; i <= dep[x]; i ++ ) {
LL v1 = 1LL, v2 = 1LL, v3 = 0;
for(auto v : E[x]) {
if(v == fa) continue;
int sz = f[v].size();
v3 = (v3 * f[v][min(sz - 1, i)] % mod + (sz - 1 >= i ? g[v][i] : 0) * v1 % mod) % mod;
v1 = v1 * f[v][min(sz - 1, i)] % mod;
v2 = v2 * (i == 0 ? 0 : f[v][min(sz - 1, i - 1)]) % mod;
}
f[x].pb((v1 - v2 + mod) % mod); // f[x][i]
if(i == 0) v3 = (v3 + 1) % mod;
g[x].pb(v3); // g[x][i]
}
}
}
int pre_max[N], suf_max[N];
vector< LL > pre_f[N], suf_f[N];
vector< LL > pre_g[N], suf_g[N];
void DP(int x, int fa) { // 有了每个点的 f 和 g, 考虑换根
// 求答案
LL ret = 0;
for(int i = 0; i < F[x].size(); i ++ ) {
F[x][i].v.pb(0); G[x][i].v.pb(0);
for(int j = F[x][i].v.size() - 1; j > 0; j -- ) F[x][i].v[j] = (F[x][i].v[j] + F[x][i].v[j - 1]) % mod;
for(int j = G[x][i].v.size() - 1; j > 0; j -- ) G[x][i].v[j] = (G[x][i].v[j] + G[x][i].v[j - 1]) % mod;
for(int j = 1; j < F[x][i].v.size(); j ++ ) F[x][i].v[j] = (F[x][i].v[j] + F[x][i].v[j - 1]) % mod; // 求出前缀和
}
for(int i = 0; i <= n / 2; i ++ ) {
LL dp[3] = {1, (i == 0), 0}; // 已经钦定了几个
for(int j = 0; j < F[x].size(); j ++ ) {
int sz = F[x][j].v.size();
LL h[3] = {0, 0, 0};
if(sz > i) { // 钦定
h[2] = (h[2] + dp[1] * G[x][j].v[i] % mod) % mod;
h[1] = (h[1] + dp[0] * G[x][j].v[i] % mod) % mod;
}
for(int k = 0; k < 3; k ++ ) h[k] = (h[k] + dp[k] * F[x][j].v[min(i, sz - 1)] % mod) % mod; // 不钦定
swap(dp, h);
}
ret = (ret + dp[2]) % mod;
}
ans = (ans + ret) % mod;
// 换根
int sz = 0; for(int i = 0; i < F[x].size(); i ++ ) sz = max(sz, (int)F[x][i].v.size());
for(int i = 0; i < F[x].size(); i ++ ) {
int ssz = F[x][i].v.size();
pre_max[i] = max(i == 0 ? 0 : pre_max[i - 1], ssz);
for(int j = 0; j < sz; j ++ ) {
pre_f[i][j] = (i == 0 ? 1 : pre_f[i - 1][j]) * F[x][i].v[min(ssz - 1, j)] % mod;
pre_g[i][j] = ((i == 0 ? 0 : pre_g[i - 1][j]) * F[x][i].v[min(ssz - 1, j)] % mod + (i == 0 ? 1 : pre_f[i - 1][j]) * (ssz - 1 >= j ? G[x][i].v[j] : 0) % mod) % mod;
}
}
for(int i = F[x].size() - 1; i >= 0; i -- ) {
int ssz = F[x][i].v.size();
suf_max[i] = max(i == F[x].size() - 1 ? 0 : suf_max[i + 1], ssz);
for(int j = 0; j < sz; j ++ ) {
suf_f[i][j] = ((i == F[x].size() - 1 ? 1 : suf_f[i + 1][j]) * F[x][i].v[min(ssz - 1, j)]) % mod;
suf_g[i][j] = ((i == F[x].size() - 1 ? 0 : suf_g[i + 1][j]) * F[x][i].v[min(ssz - 1, j)] % mod + (i == F[x].size() - 1 ? 1 : suf_f[i + 1][j]) * (ssz - 1 >= j ? G[x][i].v[j] : 0) % mod) % mod;
}
}
for(int i = 0; i < F[x].size(); i ++ ) { // 换根
int v = F[x][i].x;
if(v == fa) continue;
int ssz = max(i == 0 ? 0 : pre_max[i - 1], i == F[x].size() - 1 ? 0 : suf_max[i + 1]); // 数组大小
vector< LL > tf, tg;
if(ssz == 0) {tf.pb(1); tg.pb(1);}
else {
for(int j = 0; j < ssz; j ++ ) {
tf.pb(((i == 0 ? 1 : pre_f[i - 1][j]) * (i == F[x].size() - 1 ? 1 : suf_f[i + 1][j]) % mod - (j == 0 ? 0 : (i == 0 ? 1 : pre_f[i - 1][j - 1]) * (i == F[x].size() - 1 ? 1 : suf_f[i + 1][j - 1]) % mod) + mod) % mod);
tg.pb(((i == 0 ? 0 : pre_g[i - 1][j]) * (i == F[x].size() - 1 ? 1 : suf_f[i + 1][j]) % mod + (i == 0 ? 1 : pre_f[i - 1][j]) * (i == F[x].size() - 1 ? 0 : suf_g[i + 1][j]) % mod + (j == 0)) % mod);
}
}
F[v].pb((node) {x, tf}); G[v].pb((node) {x, tg});
}
for(int i = 0; i < F[x].size(); i ++ ) {
int v = F[x][i].x;
if(v == fa) continue;
DP(v, x);
}
}
int main() {
scanf("%d", &n);
for(int i = 1; i < n; i ++ ) {
int u, v; scanf("%d%d", &u, &v);
E[u].pb(v); E[v].pb(u);
}
for(int i = 0; i < n; i ++ ) {
for(int j = 0; j < n; j ++ ) {
pre_f[i].pb(0); pre_g[i].pb(0);
suf_f[i].pb(0); suf_g[i].pb(0);
}
}
dfs(1, 0);
DP(1, 0);
cout << ans << endl;
return 0;
}
总结
学到的比较重要的套路是：
在树的边权只有
0
/
1
0/1
0/1 时对直径计数，考虑枚举直径中点就算贡献。