【题解】P10060 [SNOI2024] 树 V 图

P10060 [SNOI2024] 树 V 图 – 洛谷 (luogu.com.cn)

我拿到了乌乌和哑哑的指偶，太萌了。

像黄豆饼和巧克力，可以亲嘴之物。

0.分析

无向树。

其实我们可以理解题意为给每个点定好了一个国家，每个国家必有一个首都。

而每个点与隶属的那个国家的首都的距离，一定是它与所有首都中最小的。

当然，与两个首都距离一样时选编号小的那个。

现在我们知道每个点隶属的国家，求有多少种不同的首都选择方案。

首先，每个国家之间一定是联通的，这里指同一个国家内不可能有另一个国家横贯，导致不连通。

不然一定可以调整为一个更合法的方案，也就是有些点距离、编号还可以更小。

那么，我们可以直接将同一个国家的连通块看成一个整体。

这样图中还需要讨论的边就是连接两个不同国家的“国境边”，

只有国境边两端的点都满足“离隶属国家首都更近（或平局按编号小）”，那么整棵树就满足要求。

为什么这么说？因为边境点没问题，内部的点就更没问题。有问题的一定先会是边境点。

1.讨论

对于国境边 (u, v)，其中 u 属于国家 x，v 属于国家 y，我们检查：

• 对于点 u：它到 x 的首都 i 的距离必须 ≤ 到 y 的首都 j 的距离 （如果 x < y，允许等号；如果 x > y，则必须 <）。

• 对于点 v：它到 y 的首都 j 的距离必须 ≤ 到 x 的首都 i 的距离 （如果 y < x，允许等号；否则必须 <）。

我们构建 vc 数组，vc[t] 表示：

在 v 联通块中，所有距离 u 恰好为 t 的那些节点，它们作为首都时的方案数之和。

对于 v 所在的连通块，我们求出它的 vc。

而对于 u 连通块的所有点，我们枚举它作为 u 连通块的首都。

这个首都成立，当且仅当 v 所在连通块的首都和 v 的距离和这个首都与 x 的距离相同。

但凡哪一方小了，都会导致至少一个点的“叛变”。

2.实现

详见代码。

#include<bits/stdc++.h>
using namespace std;

typedef long long LL;
const LL P = 998244353;
const int N = 3010;
int n, K;
vector<int> G[N], co[N];
int c[N], havec[N];
LL dp[N];
bool flag;

// 遍历 v 连通块得到 vc，x 当前节点，d 是当前节点和 v 的距离 + 1
// + 1 是为了 – 1的时候不用考虑边界，以下讨论概念都不涉及 + 1，当没有 + 1 阅读就好
// vc[t] = 在 v 联通块中，所有距离 u 恰好为 t 的那些节点，它们作为首都时的方案数之和
void get_vc(vector<LL> &vc, int x, int fa, int d) {
vc[d] = (vc[d] + dp[x]) % P;
for (int y : G[x]) if (y != fa && c[y] == c[x]) {
// 只在当前同个国家遍历
get_vc(vc, y, x, d + 1);
}
}

// 遍历 u 连通块得到 dp，x 当前节点，fa 是 v，d 是当前节点和 u 的距离 + 1
// dp[x]：当 x 作为当前连通块的首都的方案数之和（包括 x 所连接的不同的子国家）
// 因为这里的 v 对于 x 来说是完全没接触过的节点，所以计算方案数要用乘法原理
void get_dp(vector<LL> &vc, int x, int fa, int vcon, int d) {
// 如果 u 国家编号 < v 国家编号（c[x] < c[fa]）
// 那么相等时，首都 x 就更占优
// 也就是对于节点 u 来说，到 x 的距离是 d
// 对于节点 v 来说，到自己首都的距离可以是 d 或者 d – 1
// 即对于节点 u 来说，到另一个首都的距离是 d + 1 或 d
// 这不影响，因为相等有限选编号小的

// 那到 v 的首都距离能不能比 d – 1 小呢？
// 不能，不然 u 会叛变
// 那到 v 的首都距离能不能比 d 大呢？
// 不能，不然 v 会叛变

// 如果 v 国家编号 < u 国家编号（c[x] > c[fa]）
// 那么相等时，v 那边的首都就更占优
// 也就是对于节点 u 来说，到 x 的距离是 d
// 对于节点 v 来说，到自己首都的距离可以是 d 或者 d + 1
// 即对于节点 u 来说，到另一个首都的距离是 d + 2 或 d + 1
// 只有在严格比 d 大，节点 u 的首都才会选择 d

// 那到 v 的首都距离能不能比 d 小呢？
// 不能，不然 u 会叛变
// 那到 v 的首都距离能不能比 d + 1 大呢？
// 不能，不然 v 会叛变
dp[x] = dp[x] * (vc[d] + vc[d + (c[x] < vcon ? -1 : 1)] % P) % P;
for (int y : G[x]) if (y != fa && c[y] == c[x]) {
// 只在当前同个国家遍历
get_dp(vc, y, x, vcon, d + 1);
}
}

void dfs(int u, int fa) {
havec[c[u]] += (c[u] != c[fa]);
if (havec[c[u]] >= 2) {
flag = 0;
return ;
}

for (int v : G[u]) if (v != fa) {
dfs(v, u);
if (c[u] != c[v]) { // 只有国境边需要考虑更改首都
// 处理好国境边的首都分配，其他点一定没问题
vector<LL> vc(n + 5, 0);
get_vc(vc, v, u, 1);
get_dp(vc, u, v, c[v], 1);
}
}
}

int main () {
ios::sync_with_stdio(false);
cin.tie(0);

int T;
cin >> T;
while (T –) {
cin >> n >> K;
for (int i = 0; i <= n; i ++) {
c[i] = havec[i] = 0;
dp[i] = 1; // 让每个节点作为首都各有一种方案
G[i].clear();
co[i].clear();
}
for (int i = 1; i < n; i ++) {
int u, v;
cin >> u >> v;
G[u].push_back(v);
G[v].push_back(u);
}
for (int i = 1; i <= n; i ++) {
cin >> c[i];
co[c[i]].push_back(i);
}

flag = 1;
dfs(1, 0); // 我们假设 1 号节点是树根
// 所以 co[1] 就是最顶上的那个国家
// 其他国家都是它的“子国家”

for (int i = 1; i <= K; i ++) if (havec[i] != 1) {
flag = 0;
break;
}

if (flag == 0) {
cout << "0\\n";
continue;
}

LL ans = 0;
for (int i : co[c[1]]) {
// 到最后统计答案时，我们只关注 co[1] 里的点
// 因为其他点一定通过 dfs 统计贡献到这里面的点了
// 最后用加法原理将 co[1] 里每个点作为首都的方案数加在一起
ans = (ans + dp[i]) % P;
}
cout << ans << "\\n";
}

return 0;
}