找同色节点的直径端点并判共线，共线时按切断端点两侧第一条边后的连通块大小相乘计数。
OJ: luogu
题目 ID: P5588
难度:提高
标签:树的直径LCA树上计数
日期: 2026-07-16 23:59
形式化题目
给定一棵
思路
先看一个可以直接验证想法的朴素解：
cpp
/**
* Author by Rainboy blog: https://rainboylv.com github: https://github.com/rainboylvx
* rbook: -> https://rbook.roj.ac.cn https://rbook2.roj.ac.cn
* rainboy的学习导航网站: https://idx.roj.ac.cn
* create_at: 2026-08-12 22:30
* update_at: 2026-08-12 22:32
*/
// brute.cpp：小数据暴力解，用来帮助理解题意并辅助对拍。
// 直接枚举所有点对 (u,v)，收集路径上的节点集合（注意端点 u、v 本身也在路径
// 上），再检查某种颜色的全部节点是否都在这条路径上。
// 复杂度 O(n^4) 量级，只适合 n <= 30 的小数据。
#include <bits/stdc++.h>
using namespace std;
const int MAXN = 35;
int n;
int color[MAXN]; // color[i]：节点 i 的颜色
vector<int> g[MAXN]; // 邻接表存树
int parent[MAXN]; // 树上父亲（根 1 的父亲为 0）
vector<int> color_nodes[MAXN]; // color_nodes[c]：颜色 c 的全部节点
int probe[MAXN]; // probe[x]：找 LCA 时的临时探针标记（用递增 token，免清空）
int mark[MAXN]; // mark[x]：x 是否在当前的路径 u-v 上
// 从根 1 出发 BFS 求父亲。
void bfs_parent() {
int q[MAXN], head = 0, tail = 0;
q[tail++] = 1;
parent[1] = 0;
while (head < tail) {
int u = q[head++];
for (int i = 0; i < (int)g[u].size(); i++) {
int v = g[u][i];
if (v != parent[u]) {
parent[v] = u;
q[tail++] = v;
}
}
}
}
// 找 u、v 的最近公共祖先：u 一路爬到根并打探针，v 向上爬，
// 撞到的第一个探针节点就是 LCA。
int find_lca(int u, int v, int token) {
int x = u;
while (x != 0) {
probe[x] = token;
x = parent[x];
}
int y = v;
while (probe[y] != token) {
y = parent[y];
}
return y;
}
// 把路径 u-lca-v 上的全部节点（含两端点与 LCA）打上标记。
void mark_path(int u, int v, int l, int token) {
int x = u;
while (x != l) {
mark[x] = token;
x = parent[x];
}
mark[l] = token;
int y = v;
while (y != l) {
mark[y] = token;
y = parent[y];
}
}
int main() {
cin >> n;
for (int i = 1; i <= n; i++) {
cin >> color[i];
color_nodes[color[i]].push_back(i);
}
for (int i = 1; i < n; i++) {
int u, v;
cin >> u >> v;
g[u].push_back(v);
g[v].push_back(u);
}
bfs_parent();
int ans[MAXN] = {0};
int token = 0;
// 枚举所有点对 (u,v)，u < v。
for (int u = 1; u <= n; u++) {
for (int v = u + 1; v <= n; v++) {
token++;
int l = find_lca(u, v, token);
mark_path(u, v, l, token);
// 对每种颜色检查：该颜色的全部节点是否都在这条路径上。
for (int c = 1; c <= n; c++) {
bool ok = true;
for (int i = 0; i < (int)color_nodes[c].size(); i++) {
if (mark[color_nodes[c][i]] != token) {
ok = false;
break;
}
}
if (ok) {
ans[c]++;
}
}
}
}
for (int c = 1; c <= n; c++) {
cout << ans[c] << '\n';
}
return 0;
}brute.cpp 枚举所有点对
关键观察是按颜色分组看几何形状，设
：空集是任何路径的子集，答案恒为 。 ：路径必须经过单点 ，即两端点分居 的两个不同连通块。用补集算： 减去"两端点在同一连通块"的点对数。 ：同色点必须共线——全部落在某一条路径上，否则任何路径都装不下它们，答案为 。共线时，这条路径就是同色点集合的直径路径。
共线的判定与计数只需三步：
- 两次最远点扫描求直径端点：树上任意点出发的最远点必是点集直径的一个端点，所以从任意同色点出发找最远点
，再从 出发找最远点 ， 就是直径端点； - 共线检查：
在路径 上当且仅当 ，全部同色点满足才共线； - 计数：路径包含整条
，等价于两端点分别落在"切断 朝向 的第一条边"后的 侧与"切断 朝向 的第一条边"后的 侧，答案就是两侧大小之积。
以样例 1 颜色 2 为例，树是链
text
1 --- 2 --- 3 --- 4
a=2 b=3 直径路径 2-3
切断 a 朝向 b 的第一条边（边 2-3）：a 侧 = {1,2}，大小 2
切断 b 朝向 a 的第一条边（边 3-2）：b 侧 = {3,4}，大小 2
答案 = 2 x 2 = 4 点对 (1,3),(1,4),(2,3),(2,4)从图中可以看到：合法点对的第一个端点只能来自
“切断
代码
cpp
/**
* Author by Rainboy blog: https://rainboylv.com github: https://github.com/rainboylvx
* rbook: -> https://rbook.roj.ac.cn https://rbook2.roj.ac.cn
* rainboy的学习导航网站: https://idx.roj.ac.cn
* create_at: 2026-08-12 22:30
* update_at: 2026-08-12 22:31
*/
// main.cpp：对每种颜色，先找同色节点集合的直径端点（树上任意点出发的最远
// 点一定是直径端点，扫描两次即可），再判断全部同色节点是否都落在直径路径
// 上；共线时答案等于"切断直径两端各第一条边后，两个外侧连通块大小相乘"。
#include <bits/stdc++.h>
using namespace std;
const int MAXN = 1000005; // n <= 10^6
int n;
int color[MAXN]; // color[i]：节点 i 的颜色
int head[MAXN], to[MAXN * 2], nxt[MAXN * 2], edge_cnt; // 链式前向星存树
int parent[MAXN]; // 树上父亲（根 1 的父亲为 0）
int depth[MAXN]; // 节点深度（根为 0）
int order[MAXN]; // 从根出发的 BFS 顺序
int subtree[MAXN]; // subtree[u]：以 u 为根的子树大小
vector<vector<int>> up; // up[k][x]：x 向上跳 2^k 步到达的祖先
int color_head[MAXN]; // color_head[c]：颜色 c 的链表头节点
int next_same[MAXN]; // next_same[x]：与 x 同色的下一个节点
int color_count[MAXN]; // 每种颜色的节点个数
long long ans[MAXN]; // 每种颜色的答案
// 加一条无向边。
void add_edge(int u, int v) {
edge_cnt++;
to[edge_cnt] = v;
nxt[edge_cnt] = head[u];
head[u] = edge_cnt;
}
// 快读：n 可达 10^6，用 getchar 手写读整数更稳。
int read_int() {
int x = 0;
char c = getchar();
while (c < '0' || c > '9') {
c = getchar();
}
while (c >= '0' && c <= '9') {
x = x * 10 + c - '0';
c = getchar();
}
return x;
}
// 从根 1 出发 BFS，求出父亲、深度和 BFS 顺序（迭代，避免百万深度递归爆栈）。
void bfs_build() {
int tail = 0;
order[tail++] = 1;
parent[1] = 0;
depth[1] = 0;
for (int i = 0; i < tail; i++) {
int u = order[i];
for (int e = head[u]; e != 0; e = nxt[e]) {
int v = to[e];
if (v != parent[u]) {
parent[v] = u;
depth[v] = depth[u] + 1;
order[tail++] = v;
}
}
}
}
// 节点 x 向上跳 steps 步。
int jump(int x, int steps) {
for (int k = 0; steps > 0; k++) {
if (steps & 1) {
x = up[k][x];
}
steps >>= 1;
}
return x;
}
// 求 x、y 的最近公共祖先（倍增 LCA）。
int lca(int x, int y) {
if (depth[x] < depth[y]) {
swap(x, y);
}
x = jump(x, depth[x] - depth[y]);
if (x == y) {
return x;
}
for (int k = (int)up.size() - 1; k >= 0; k--) {
if (up[k][x] != up[k][y]) {
x = up[k][x];
y = up[k][y];
}
}
return parent[x];
}
// 树上两点距离。
int dist(int x, int y) {
int a = lca(x, y);
return depth[x] + depth[y] - 2 * depth[a];
}
// 切断从 x 朝向 y 的第一条边后，x 所在连通块的大小。
long long side_size(int x, int y) {
int a = lca(x, y);
if (a == x) {
// y 在 x 的子树里：把 y 上跳到 x 的直接儿子 child。
int child = jump(y, depth[y] - depth[x] - 1);
return n - subtree[child]; // x 那一侧 = 整棵树去掉 child 的子树
}
return subtree[x];
}
int main() {
n = read_int();
for (int i = 1; i <= n; i++) {
color[i] = read_int();
}
for (int i = 1; i < n; i++) {
int u = read_int();
int v = read_int();
add_edge(u, v);
add_edge(v, u);
}
bfs_build();
// 子树大小：所有叶子先为 1，按 BFS 逆序累加到父亲（深度大的先算）。
for (int i = 1; i <= n; i++) {
subtree[i] = 1;
}
for (int i = n - 1; i >= 1; i--) {
subtree[parent[order[i]]] += subtree[order[i]];
}
// 倍增祖先表：up[k][x] = up[k-1][ up[k-1][x] ]。
int LOG = 0;
while ((1 << LOG) <= n) {
LOG++;
}
up.resize(LOG);
for (int k = 0; k < LOG; k++) {
up[k].resize(n + 1);
}
for (int x = 1; x <= n; x++) {
up[0][x] = parent[x];
}
for (int k = 1; k < LOG; k++) {
for (int x = 1; x <= n; x++) {
up[k][x] = up[k - 1][up[k - 1][x]];
}
}
// 相同颜色的节点建成链表：color_head[c] -> next_same。
for (int x = 1; x <= n; x++) {
int c = color[x];
next_same[x] = color_head[c];
color_head[c] = x;
color_count[c]++;
}
long long all_pairs = 1LL * n * (n - 1) / 2; // 全部无序点对
for (int c = 1; c <= n; c++) {
ans[c] = all_pairs; // 没有该颜色节点时，所有点对都合法
}
for (int c = 1; c <= n; c++) {
if (color_count[c] == 0) {
continue;
}
if (color_count[c] == 1) {
// 只有一个同色节点 x：路径必须经过 x，即两端点落在 T - x 的
// 不同连通块里。用补集：总数减去两端点在同一连通块的点对。
int x = color_head[c];
long long excluded = 1LL * (n - subtree[x]) * (n - subtree[x] - 1) / 2;
for (int e = head[x]; e != 0; e = nxt[e]) {
int v = to[e];
if (parent[v] == x) {
long long s = subtree[v];
excluded += s * (s - 1) / 2;
}
}
ans[c] = all_pairs - excluded;
continue;
}
// 多个同色节点：两次最远点扫描求出同色点集合的直径端点。
int first = color_head[c];
int second = first;
int best = -1;
for (int x = first; x != 0; x = next_same[x]) {
int d = dist(first, x);
if (d > best) {
best = d;
second = x;
}
}
int third = second;
best = -1;
for (int x = first; x != 0; x = next_same[x]) {
int d = dist(second, x);
if (d > best) {
best = d;
third = x;
}
}
int dia = dist(second, third);
// 检查全部同色节点是否都在直径路径 second-third 上。
// 在路径上当且仅当 dist(second,x) + dist(x,third) == dia。
bool lies_on_path = true;
for (int x = first; x != 0; x = next_same[x]) {
if (dist(second, x) + dist(x, third) != dia) {
lies_on_path = false;
break;
}
}
if (lies_on_path) {
// 路径要包含整条 second-third，等价于两端点分居两个外侧连通块。
ans[c] = side_size(second, third) * side_size(third, second);
} else {
ans[c] = 0; // 同色节点不共线，任何一条路径都装不下
}
}
for (int c = 1; c <= n; c++) {
printf("%lld\n", ans[c]);
}
return 0;
}复杂度
- 时间：
。建树与子树 ，倍增表 ；每种颜色扫 3 遍同色链表，共 次 的距离查询。 - 空间：倍增表
，其余数组 。
总结
"路径必须覆盖一组点"的通用套路是三步：求这组点的直径端点、检查全部点是否落在直径路径上（共线）、共线后把计数化为切断直径两端第一条边后的连通块大小相乘。核心工具是树上距离（倍增 LCA）与子树大小。rbook 的《倍增求 LCA》讲解了本解使用的 jump / lca / dist 模板（lca-binary-lifting），本解在其基础上为 side_size 计数。
图示解析
这张 ASCII 图展示整道题的解题路线：
text
朴素枚举（brute.cpp）
枚举点对 (u,v)，收集路径节点，检查每种颜色的全部节点是否都在路径上
复杂度 O(n^4)，只适合 n <= 30
|
| 瓶颈：枚举所有点对 + 逐对收集路径，结果只由颜色节点集合的形状决定
v
关键观察：按颜色分组看几何形状
|S|=0：答案 = C(n,2)
|S|=1：补集，C(n,2) - 各连通块内部点对数
|S|>=2：同色点集合必须共线，否则任何路径都装不下，答案为 0
|
v
共线判定与计数（main.cpp）
两次最远点扫描 -> 直径端点 a,b（任意点出发的最远点必是直径端点）
共线检查：dist(a,x) + dist(x,b) == dist(a,b) 对所有同色点成立
计数：切断 a、b 朝向对方的第一条边，两侧连通块大小相乘
距离与"切断后的侧大小"都由倍增 LCA + subtree 在 O(log n) 得到
|
v
复杂度 O(n log n)，空间 O(n log n)图中四条主线对应"暴力慢在哪"“按颜色分类后各是什么形状”“共线如何判定与计数”“正式解如何复用 LCA 与子树大小”。核心是把"路径必须包含一组点"翻译成"这组点共线 + 两端点分居直径路径两侧"，计数就退化成两个连通块大小的乘积。
