小猪佩奇爬树

形式化题目

给定一棵 $n$ 个节点的树, 每个节点有一个颜色 $w_i \in [1,n]$ . 对每种颜色 $c$ , 统计无序点对 $\{u,v\}$ ( $u \neq v$ ) 的数量, 使 $u$ 到 $v$ 的唯一树上路径经过的节点集合包含全部颜色为 $c$ 的节点. 端点 $u$ 、 $v$ 本身也算在路径上, 其他颜色的节点经过与否不限.

记 $S_c$ 为颜色 $c$ 的节点集合, $P(u,v)$ 为路径 $u \to v$ 上的节点集合, 则要求的是满足 $S_c \subseteq P(u,v)$ 的点对数量.

解法总览

解法二和解法三按"如何刻画同色点的几何形状"区分; 解法一只是对拍与理解题意的基线, 不参与这层刻画.

解法	核心观察	复杂度	适用场景	对应代码
一、暴力枚举	直接枚举点对并收集路径节点集合	$O(n^3)$	只作对拍基线, 不参与形状分类	`brute.cpp`
二、倍增 LCA + 直径端点	共线 $\iff$ 全部同色点在集合直径路径上, 答案 = 切断两端第一条边后两侧大小之积	$O(n\log n)$ 时间, $O(n\log n)$ 空间	复用已有的 LCA 与子树大小工具, 思路几何直观	`main2.cpp`
三、数方向	在节点 $u$ 处看同色点分布在几个方向: 三个方向判死, 只有一个方向则 $u$ 是端点	$O(n)$ 时间, $O(n)$ 空间	正式主解, 不需要 LCA 与倍增表	`main.cpp` / `main.py`

解法三是正式主解. 它把"同色点是否共线"这个几何判断换成一个纯计数条件, 因此既不需要求距离, 也不需要倍增表, 时间和空间都少一个 $\log$ 因子. 解法二的价值在于推导自然: 读者已经会用 LCA 求距离, 从"路径必须包含点集"想到"点集的直径路径"只差一步. 两条路线互不依赖, 所以并列展示, 不构成前后优化关系.

两个正解共用同一层入口: 先按 $|S_c|$ 分成三类, 只有最后一类需要真正的算法.

$|S_c| = 0$ : 空集含于任何集合, 所有 $\binom{n}{2}$ 个点对都合法.
$|S_c| = 1$ , 设唯一同色点为 $x$ : 路径必须经过 $x$ , 等价于两端点落在 $T - x$ 的两个不同连通块里. 用补集算: 总数减去"两端点在同一连通块"的点对数.
$|S_c| \geqslant 2$ : 一条路径能否同时装下若干点, 取决于这些点是否落在一条链上. 若 $S_c$ 的最小包含子树存在分叉, 任何路径都装不下, 答案为 0.

所以真正的难点只有第三类: 判断 $S_c$ 是否共线, 并在共线时数出合法点对.

解法一: 暴力枚举

思路

先写下能直接验证题意的朴素解, 后面所有优化都以它为标准答案.

枚举所有点对 $(u,v)$ , 收集路径上的全部节点, 再对每种颜色检查它的全部节点是否都在这条路径上. 收集路径的做法是: 先求出 $u$ 、 $v$ 的 LCA, 然后从 $u$ 、 $v$ 分别向上爬到 LCA, 沿途打标记. 注意两端点本身也在路径上.

枚举量是 $O(n^2)$ , 每对点收集路径 $O(n)$ . 检查颜色那一层看着是" $n$ 种颜色各 $O(n)$ ", 但各颜色的节点数之和恒为 $n$ , 所以每对点实际只花 $O(n)$ . 总复杂度 $O(n^3)$ . 代码里 MAXN = 35, 所以它只用于 $n \leqslant 30$ 的对拍与理解题意.

代码

cpp

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117/**
 * 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-10-26 10:23
 */
// brute.cpp：小数据暴力解，用来帮助理解题意并辅助对拍。
// 直接枚举所有点对 (u,v)，收集路径上的节点集合（注意端点 u、v 本身也在路径
// 上），再检查某种颜色的全部节点是否都在这条路径上。
// 复杂度 O(n^3)：点对 O(n^2)，每对收集路径 O(n)；检查颜色时各颜色的节点数
// 之和恒为 n，所以也是对点 O(n)。只适合 n <= 30 的小数据（数组上界 MAXN）。
#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;
}

复杂度

时间 $O(n^3)$ , 空间 $O(n)$ . 只用于对拍.

解法二: 倍增 LCA + 直径端点

思路

要判断 $S_c$ 是否共线, 只需要看这个点集的直径. 树上有一个常用性质: 从任意点出发找最远点, 这个最远点一定是点集直径的一个端点. 所以从 $S_c$ 中任取一点出发找最远点 $a$ , 再从 $a$ 出发找最远点 $b$ , $(a,b)$ 就是 $S_c$ 的直径端点.

有了 $a$ 、 $b$ , 判定和计数都用距离表达:

共线检查: $x$ 在路径 $a \to b$ 上, 当且仅当 $\operatorname{dist}(a,x) + \operatorname{dist}(x,b) = \operatorname{dist}(a,b)$ . 三角不等式说明左边恒 $\geqslant$ 右边, 取等表示三条路径在 $x$ 处合并. 全部同色点都满足才共线.
计数: 路径要包含整条 $a \to b$ , 等价于两端点分别落在"切断 $a$ 朝向 $b$ 的第一条边"后的 $a$ 侧、"切断 $b$ 朝向 $a$ 的第一条边"后的 $b$ 侧. 两侧独立, 答案为两侧大小之积.

下图的树是样例 1 的链 $1\text{-}2\text{-}3\text{-}4$ , 取颜色 2, $S_2 = \{2,3\}$ . 图中把两个直径端点 $a=2$ 、 $b=3$ 上色, 并标出切断边:

flowchart LR
  N1["1"] --- N2["2 (a)"]
  N2 -.->|"切断 2-3"| N3["3 (b)"]
  N3 --- N4["4"]

  style N2 fill:#cfe8ff,stroke:#3b82f6
  style N3 fill:#cfe8ff,stroke:#3b82f6

这里有一个容易被忽略的边界情形: $a$ 、 $b$ 在树上相邻, 所以" $a$ 朝向 $b$ 的第一条边"和" $b$ 朝向 $a$ 的第一条边"其实是同一条边 $2\text{-}3$ . 切掉它以后, $a$ 侧的连通块是 $\{1,2\}$ 大小 2, $b$ 侧的连通块是 $\{3,4\}$ 大小 2, 两侧仍然互不重叠, 公式照常成立.

合法点对的第一个端点只能来自 $a$ 侧, 第二个只能来自 $b$ 侧, 两侧互相独立, 所以直接相乘. 这与题面样例说明列出的 4 个点对一致.

最后需要"切断 $x$ 朝向 $y$ 的第一条边后 $x$ 侧的大小" $\operatorname{side}(x,y)$ , 用 LCA 和子树大小求:

若 $\operatorname{lca}(x,y) = x$ , 说明 $y$ 在 $x$ 子树内. 把 $y$ 向上跳 $\operatorname{depth}[y] - \operatorname{depth}[x] - 1$ 步得到 $x$ 通往 $y$ 的直接儿子 $t$ , 则 $x$ 侧 = $n - \operatorname{subtree}[t]$ .
否则 $y$ 不在 $x$ 的子树里, $x$ 侧就是 $\operatorname{subtree}[x]$ .

树上距离都由倍增 LCA 求出, 每次 $O(\log n)$ .

代码

cpp

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236/**
 * 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-10-26 17:15
 */
// main2.cpp：对每种颜色，先找同色节点集合的直径端点（树上任意点出发的最远
// 点一定是直径端点，扫描两次即可），再判断全部同色节点是否都落在直径路径
// 上；共线时答案等于"切断直径两端各第一条边后，两个外侧连通块大小相乘"。
#include <bits/stdc++.h>
using namespace std;

const int MAXN = 1000005; // n <= 10^6

int n;
int color[MAXN]; // color[i]：节点 i 的颜色

vector<int> g[MAXN]; // g[u]：与 u 相邻的所有点

int parent[MAXN]; // 树上父亲（根 1 的父亲为 0）
int depth[MAXN];  // 节点深度（根为 0）
int order[MAXN];  // 从根出发的 BFS 顺序
int subtree[MAXN]; // subtree[u]：以 u 为根的子树大小

const int MAXLOG = 21; // 2^20 > 10^6，倍增层数上限
int up[MAXN][MAXLOG];  // up[x][k]：x 向上跳 2^k 步到达的祖先
int LOG;               // 实际用到的倍增层数，满足 2^LOG > n

vector<int> color_nodes[MAXN]; // color_nodes[c]：颜色 c 的所有节点
int color_count[MAXN];         // 每种颜色的节点个数

long long ans[MAXN];  // 每种颜色的答案

// 快读：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 v : g[u]) {
            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[x][k];
        }
        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 = LOG - 1; k >= 0; k--) {
        if (up[x][k] != up[y][k]) {
            x = up[x][k];
            y = up[y][k];
        }
    }
    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];
}

// 求颜色 c 的节点集合的直径端点。要求该颜色至少有一个节点。
// 两次最远点扫描：从任意同色点出发找到最远点 a，再从 a 找到最远点 b，(a, b) 即直径端点。
pair<int, int> color_diameter(int c) {
    int start = color_nodes[c][0];
    int far1 = start;
    int best = -1;
    for (int x : color_nodes[c]) {
        int d = dist(start, x);
        if (d > best) {
            best = d;
            far1 = x;
        }
    }
    int far2 = far1;
    best = -1;
    for (int x : color_nodes[c]) {
        int d = dist(far1, x);
        if (d > best) {
            best = d;
            far2 = x;
        }
    }
    return make_pair(far1, far2);
}

// 判断颜色 c 的所有节点是否都落在路径 a-b 上。
// 点在路径上当且仅当 dist(a,x) + dist(x,b) == dist(a,b)，这里 dia 就是该值。
bool all_on_path(int c, int a, int b, int dia) {
    for (int x : color_nodes[c]) {
        if (dist(a, x) + dist(x, b) != dia) {
            return false;
        }
    }
    return true;
}

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();
        g[u].push_back(v);
        g[v].push_back(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[x][k] = up[ up[x][k-1] ][k-1]。
    LOG = 0;
    while ((1 << LOG) <= n) {
        LOG++;
    }
    for (int x = 1; x <= n; x++) {
        up[x][0] = parent[x];
    }
    for (int k = 1; k < LOG; k++) {
        for (int x = 1; x <= n; x++) {
            up[x][k] = up[up[x][k - 1]][k - 1];
        }
    }

    // 按颜色把节点分组：color_nodes[c] 存颜色 c 的所有节点。
    for (int x = 1; x <= n; x++) {
        int c = color[x];
        color_nodes[c].push_back(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_nodes[c][0];
            long long excluded = 1LL * (n - subtree[x]) * (n - subtree[x] - 1) / 2;
            for (int v : g[x]) {
                if (parent[v] == x) {
                    long long s = subtree[v];
                    excluded += s * (s - 1) / 2;
                }
            }
            ans[c] = all_pairs - excluded;
            continue;
        }

        // 多个同色节点：先求同色点集合的直径端点，再检查它们是否都在直径路径上。
        pair<int, int> ends = color_diameter(c);
        int a = ends.first;
        int b = ends.second;
        int dia = dist(a, b);

        if (all_on_path(c, a, b, dia)) {
            // 路径要包含整条 a-b，等价于两端点分居两个外侧连通块。
            ans[c] = side_size(a, b) * side_size(b, a);
        } else {
            ans[c] = 0; // 同色节点不共线，任何一条路径都装不下
        }
    }

    for (int c = 1; c <= n; c++) {
        printf("%lld\n", ans[c]);
    }

    return 0;
}

复杂度

时间 $O(n\log n)$ : 建树、子树大小、倍增表, 加每种颜色的两次最远点扫描与一次路径检查, 每次距离查询 $O(\log n)$ . 空间 $O(n\log n)$ : 倍增表约 $n\log n$ 个 int, $n = 10^6$ 时约 80MB; 邻接表与按颜色分组的 vector 各有 $n$ 个空容器, 每个约 24 字节, 合计约 48MB.

实测 $n = 10^6$ (判题环境, 限制 3s / 500MB): 双色链 1576ms, 单色链 653ms, 菊花 1059ms, 随机树 1663ms, 完全二叉树 1659ms, 峰值内存 208–219MB. 都在限内, 最慢一处留约 1.3s 余量.

倍增表按 up[x][k] 存放, 即"一个节点一行"的连续 $N \times 21$ 二维数组. 同样满足 $O(n\log n)$ 空间, 但比"按层分开的 $\log n$ 个一维数组"少一层指针跳转, 距离查询访问相邻内存即可; 实测这一存放顺序在链式数据上快约一倍.

解法三: 数方向

关键观察

解法二把"共线"表达成了距离等式, 代价是每次判断都要一次 LCA. 换个角度看: 共线与否, 本质上是同色点在各节点处朝几个方向散开.

任取一个根. 对颜色 $c$ 和它的一个节点 $u$ , 把 $S_c \setminus \{u\}$ 按相对 $u$ 的位置分成两部分:

儿子方向: 某个儿子 $v$ 的子树里含 $c$ 色节点. 这样的儿子可能有多个.
父亲方向: $u$ 的子树之外还有 $c$ 色节点, 即 $|S_c| - \text{inside}_u > 0$ , 其中 $\text{inside}_u$ 是 $u$ 子树内(含 $u$ )的 $c$ 色节点数.

记 $positive_u$ 为"子树里含 $c$ 色节点"的儿子个数, 定义方向数

dirs(u) = positive_u + [\,|S_c| - \text{inside}_u > 0\,]

下图画出同一个颜色在三个不同位置上的方向数, 对应要区分的三种情形:

flowchart TB
  subgraph A["情形 A: dirs = 3, 答案为 0"]
    A0["同色点"] --- AU["u"]
    AU --- A1["同色点"]
    AU --- A2["同色点"]
    AU --- A3["同色点"]
  end
  subgraph B["情形 B: dirs = 1, u 是端点"]
    BU["u"] --- B1["同色点"]
    B1 --- B2["同色点"]
  end
  subgraph C["情形 C: dirs = 2, u 在中间"]
    C1["同色点"] --- CU["u"]
    CU --- C2["同色点"]
  end

图中每条边代表一个"方向": 同色点画在边上, 表示它落在该方向的子树里. 情形 A 中 $u$ 朝三个方向都有同色点, 最小包含子树在 $u$ 处分叉, 不是一条路径(方向再多也只会更糟). 情形 B 中同色点只在 $u$ 的一侧延伸, $u$ 就是最外层的端点. 情形 C 中 $u$ 两侧都有同色点, 它只是路径内部的一个点.

三种情形的含义:

$dirs(u) \geqslant 3$ : 同色点分处三个互不相同的方向, 最小包含子树在 $u$ 处分叉, 一定不是一条路径. 此时颜色 $c$ 的答案为 0.
$dirs(u) = 1$ : 所有同色点都在 $u$ 的一侧. $u$ 只可能朝一个方向延伸, 所以它是同色点集合最外层的点, 也就是共线情形下的端点.
$dirs(u) = 2$ : $u$ 两侧都有同色点, 是路径内部点, 不是端点.

当 $|S_c| \geqslant 2$ 时 $dirs(u)$ 恒 $\geqslant 1$ : 另一个同色点要么在 $u$ 的某棵子树里, 要么在 $u$ 的子树外. 所以这里不会出现 $dirs(u) = 0$ ; 那个值只在 $|S_c| = 1$ 时出现, 而那种情况由补集分支单独处理.

这里需要一个能直接对应到"共线"的结论. 记 $T'$ 为 $S_c$ 的最小包含子树. $T'$ 的叶子一定属于 $S_c$ : 否则删掉这个叶子后剩下的子树仍包含 $S_c$ , 与最小性矛盾. 于是有

dirs(u) = 1 \iff u \text{ 是 } T' \text{ 的叶子}

两个方向各证一次:

若 $u$ 是 $T'$ 的叶子, 删掉 $u$ 后剩下的同色点仍在 $T'$ 里且连通, 所以它们全部落在 $u$ 的同一个方向, 得 $dirs(u) = 1$ .
若 $dirs(u) = 1$ , 说明 $S_c \setminus \{u\}$ 全在 $u$ 的同一个方向. 于是 $T'$ 就是把 $u$ 接到 $S_c \setminus \{u\}$ 的最小包含子树上, 接点只有一条边通向 $u$ , 所以 $u$ 在 $T'$ 中的度数为 1, 是叶子.

而一棵树是路径当且仅当它恰好有两个叶子. 再加上" $S_c$ 落在一条链上 $\Rightarrow$ $T'$ 含于该链 $\Rightarrow$ $T'$ 是路径", 两个方向都通了, 于是

S_c \text{ 共线} \iff leaf\_cnt_c = 2

其中 $leaf\_cnt_c$ 是扫描中数到的 $dirs = 1$ 的点个数. 这一步把"共线"这个几何条件完全换成了计数条件, 不再需要任何距离查询.

代码里另外还维护了一个标记 too_many[c]: 一旦发现某个同色点 $dirs \geqslant 3$ , 说明 $T'$ 在该点处分叉, 沿三个方向各走到叶子就至少有三个叶子, 不可能恰好两个. 注意它只是额外标记, 并不会中断扫描; 最终判据仍然是 $leaf\_cnt_c = 2$ .

端点 $u$ 的外侧连通块大小, 按同色点的位置分两种:

同色点全在某个儿子 only_child 的子树里( $positive_u = 1$ ): $u$ 这一侧是整棵树去掉那棵子树, 大小 $n - \operatorname{subtree}[\text{only\_child}]$ .
同色点全在 $u$ 的子树外( $positive_u = 0$ ): $u$ 这一侧就是自己的子树, 大小 $\operatorname{subtree}[u]$ .

思路

实现只需要一次自底向上的扫描:

从根出发做一次遍历, 进入节点时记下父亲, 等儿子全部处理完再结算该节点, 天然就是自底向上.
对每个 $u$ 算出 $\operatorname{subtree}[u]$ 、 $\text{inside}_u$ 和 $positive_u$ . 若 $dirs(u) \geqslant 3$ 就给该颜色打上 too_many; 若 $dirs(u) = 1$ 就把 $u$ 记作端点并算出它那一侧的大小.
最后按颜色输出: 无点给 $\binom{n}{2}$ , 单点用补集算, 多点时端点恰好两个就相乘, 否则为 0.

判断"某个儿子的子树里是否含父颜色节点"有个省事的写法. 维护 done_cnt[c] 表示颜色 $c$ 中已经处理完(子树扫完)的节点数. 父节点在进入某个儿子之前记下快照, 该儿子返回后再看 done_cnt 是否变化, 就知道这棵子树里有没有同色点. 这样不必额外为每对(颜色, 子树)存计数.

这里有个实现上的坑: 快照必须在进入儿子之前取. 如果取在儿子返回之后, 差值恒为 0, 所有儿子的方向都会漏判.

子树大小也是同一个套路: 每帧用 sub 字段累加, 儿子结算后把它加到父亲帧里(p.sub += sub), 父节点自己结算时才写进 subtree[u]. 所以 Frame 里的 sub 既是帧字段又是最终的子树大小.

树可能退化成深度 $10^6$ 的链. 递归 DFS 在这种输入上实测需要约 128MB 栈(64MB 会 signal 11), 所以这一版按仓库约定改成显式栈迭代, 把深度从栈上移走.

需要说明的是, roj-local-judge 的栈上限默认等于内存上限(本题 500MB), 递归版在这个口径下同样能过; 显式栈的收益是不再依赖判题器的栈配置, 也能在内存紧张时少占一份调用栈. 两种写法都对, 这里选迭代是为了让性能结论不挂在一个隐含前提上.

代码

C++ 版(正式主解):

cpp

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213/**
 * 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-10-26 09:20
 * update_at: 2026-10-26 17:02
 */
// main.cpp：O(n) 线性解法。自底向上扫一遍，对每个节点 u 看与 u 同色的其他节点
// 分布在几个方向上：含同色点的儿子子树各算一个方向，u 的子树外若还有同色点
// 再算一个。方向数 >= 3 说明同色点不共线，答案 0；恰好 1 说明 u 是同色点集合
// 最外层的端点，一个颜色最多两个端点，两端「外侧连通块大小」相乘即答案。
// 树可能退化成深度 10^6 的链，递归必然爆栈，所以按仓库规矩改为显式栈迭代 DFS。
#include <bits/stdc++.h>
using namespace std;

const int MAXN = 1000005; // n <= 10^6

int n;
int color[MAXN]; // color[i]：节点 i 的颜色

vector<int> g[MAXN]; // g[u]：与 u 相邻的所有点

int parent[MAXN];  // parent[x]：树上父亲（根 1 的父亲为 0）
int subtree[MAXN]; // subtree[x]：以 x 为根的子树大小

int total_cnt[MAXN];  // total_cnt[c]：颜色 c 的节点总数
int first_node[MAXN]; // first_node[c]：颜色 c 出现的第一个节点，供单点补集计算使用
int done_cnt[MAXN];   // done_cnt[c]：颜色 c 中已经处理完（子树扫完）的节点个数

int leaf_cnt[MAXN];          // leaf_cnt[c]：颜色 c 已经找到的端点个数
int too_many[MAXN];          // too_many[c]：颜色 c 的分叉方向 >= 3，答案必为 0
long long side_first[MAXN];  // side_first[c]：颜色 c 第一个端点外侧连通块的大小
long long side_second[MAXN]; // side_second[c]：颜色 c 第二个端点外侧连通块的大小

long long answer[MAXN]; // answer[c]：颜色 c 的答案

// 迭代 DFS 的一帧，记录"正在处理节点 u 及其儿子"这件事进行到哪了。
struct Frame {
    int u;          // 当前节点
    int idx;        // 下一条待处理的邻接下标（等于 g[u] 的长度表示儿子已全部处理完）
    long long sub;  // subtree[u] 的累加值
    int same;       // 子树中颜色为 color[u] 的节点数（含 u 自己）
    int positive;   // 已处理完的儿子中，子树里含 color[u] 的个数
    int only_child; // 最近一个含 color[u] 的儿子（方向数为 1 时用它算外侧大小）
    int before;     // 进入当前儿子之前 done_cnt[color[u]] 的快照
};
Frame stk[MAXN]; // 手工递归栈，最深可达 n 层，所以开成全局数组

// 快读：n 可达 10^6，输入约 3*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;
}

// 压入节点 u 的一帧，开始处理它的儿子。
void push_frame(int& top, int u) {
    top++;
    stk[top].u = u;
    stk[top].idx = 0;
    stk[top].sub = 1;
    stk[top].same = 1;
    stk[top].positive = 0;
    stk[top].only_child = 0;
    stk[top].before = 0;
}

// u 的子树全部处理完时调用：判断 u 是不是颜色 color[u] 的端点。
// inside 是 u 子树内（含 u）的同色节点数，positive 是含同色节点的儿子个数。
void finish_node(int u, int inside, int positive, int only_child) {
    int c = color[u];
    if (total_cnt[c] == 1) {
        return; // 只有 1 个同色节点，答案按补集单独算
    }

    // 同色点分布在几个方向：每个含同色点的儿子算一个方向，
    // u 的子树外若还有同色点，再算一个方向。
    int outside = total_cnt[c] - inside;
    int dirs = positive + (outside > 0 ? 1 : 0);

    if (dirs >= 3) {
        too_many[c] = 1; // 同色点朝三个方向散开，不存在能装下它们的路径
        return;
    }
    if (dirs != 1) {
        return; // dirs == 2：u 夹在中间，是同色点集合的内部节点，不是端点
    }

    // dirs == 1：同色点全在 u 的一侧，u 就是最外层的端点。
    // 外侧连通块大小 = 切掉"含同色点的那一侧"之后，u 这侧剩下的点数。
    long long side;
    if (positive == 1) {
        // 同色点全在儿子 only_child 的子树里：u 这侧 = 整棵树去掉 only_child 的子树。
        side = 1LL * n - subtree[only_child];
    } else {
        // 同色点全在 u 子树外：u 这侧就是自己的整棵子树。
        side = subtree[u];
    }
    leaf_cnt[c]++;
    if (leaf_cnt[c] == 1) {
        side_first[c] = side;
    } else if (leaf_cnt[c] == 2) {
        side_second[c] = side;
    }
}

// 显式栈迭代 DFS，等价于递归的"进入 u → 依次处理儿子 → 离开 u"。
// 离开 u 时儿子都已算完，可以确定 subtree[u] 并判断端点。
void dfs_all() {
    int top = -1;
    parent[1] = 0;
    push_frame(top, 1);

    while (top >= 0) {
        Frame &cur = stk[top];
        if (cur.idx < (int)g[cur.u].size()) {
            // 还有儿子没处理，取一个相邻点。
            int v = g[cur.u][cur.idx];
            cur.idx++;
            if (v == parent[cur.u]) {
                continue; // 跳过父亲这条边（存的是无向边）
            }
            parent[v] = cur.u;
            cur.before = done_cnt[color[cur.u]]; // 进入儿子前的快照
            push_frame(top, v);                  // 先进入儿子，父亲这一帧留在栈里
            continue;
        }

        // 儿子都处理完了，可以结算 u。
        int u = cur.u;
        int inside = cur.same;
        int positive = cur.positive;
        int only_child = cur.only_child;
        long long sub = cur.sub;
        subtree[u] = sub;
        done_cnt[color[u]]++;
        top--;

        finish_node(u, inside, positive, only_child);

        if (top >= 0) {
            // 把 u 这棵子树的结果累加到父亲身上。
            Frame &p = stk[top];
            int parent_color = color[p.u];
            int added = done_cnt[parent_color] - p.before; // 这棵子树里的父颜色节点数
            p.sub += sub;
            p.same += added;
            if (added > 0) {
                p.positive++;
                p.only_child = u;
            }
        }
    }
}

int main() {
    n = read_int();
    for (int i = 1; i <= n; i++) {
        color[i] = read_int();
        total_cnt[color[i]]++;
        if (first_node[color[i]] == 0) {
            first_node[color[i]] = i;
        }
    }
    for (int i = 1; i < n; i++) {
        int u = read_int();
        int v = read_int();
        g[u].push_back(v);
        g[v].push_back(u);
    }

    long long all_pairs = 1LL * n * (n - 1) / 2; // 全部无序点对
    for (int c = 1; c <= n; c++) {
        answer[c] = all_pairs; // 没有该颜色节点时，所有点对都合法
    }

    dfs_all();

    for (int c = 1; c <= n; c++) {
        if (total_cnt[c] == 0) {
            continue;
        }
        if (total_cnt[c] == 1) {
            // 只有一个同色节点 x：路径必须经过 x，用补集减去"两端在同一连通块"。
            int x = first_node[c];
            long long out = n - subtree[x];
            long long excluded = out * (out - 1) / 2;
            for (int v : g[x]) {
                if (parent[v] == x) {
                    long long s = subtree[v];
                    excluded += s * (s - 1) / 2;
                }
            }
            answer[c] = all_pairs - excluded;
        } else {
            // 恰好两个端点才共线；端点不是两个说明同色点不共线或分叉，答案为 0。
            answer[c] = (too_many[c] == 1 || leaf_cnt[c] != 2) ? 0 : side_first[c] * side_second[c];
        }
    }

    for (int c = 1; c <= n; c++) {
        printf("%lld\n", answer[c]);
    }

    return 0;
}

Python 版(同一算法, 适用范围见下面的复杂度一节):

python

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147import sys
from array import array

data = sys.stdin.buffer.read().split()
n = int(data[0])
color = array("i", [0])
color.extend(map(int, data[1 : n + 1]))

# 链式前向星存树：head[u] 是 u 的第一条出边，to/nxt 用 -1 表示空。
head = array("i", [-1]) * (n + 1)
to = array("i", [0]) * (2 * n - 2)
nxt = array("i", [0]) * (2 * n - 2)
edge_cnt = 0
pos = n + 1
for _ in range(n - 1):
    u = int(data[pos])
    v = int(data[pos + 1])
    pos += 2
    to[edge_cnt] = v
    nxt[edge_cnt] = head[u]
    head[u] = edge_cnt
    edge_cnt += 1
    to[edge_cnt] = u
    nxt[edge_cnt] = head[v]
    head[v] = edge_cnt
    edge_cnt += 1

total_cnt = array("i", [0]) * (n + 1)
first_node = array("i", [0]) * (n + 1)
for i in range(1, n + 1):
    c = color[i]
    total_cnt[c] += 1
    if first_node[c] == 0:
        first_node[c] = i

parent = array("i", [0]) * (n + 1)
subtree = array("i", [0]) * (n + 1)
done_cnt = array("i", [0]) * (n + 1)       # done_cnt[c]：颜色 c 已处理完的节点数
leaf_cnt = array("i", [0]) * (n + 1)       # leaf_cnt[c]：颜色 c 已找到的端点个数
too_many = array("i", [0]) * (n + 1)       # too_many[c]：分叉方向 >= 3，答案 0
side_first = array("q", [0]) * (n + 1)     # 第一个端点的外侧连通块大小
side_second = array("q", [0]) * (n + 1)    # 第二个端点的外侧连通块大小
answer = array("q", [0]) * (n + 1)

# 显式栈迭代 DFS：等价于递归「进入 u -> 依次处理儿子 -> 离开 u」。
# 树可能退化成深度 10^6 的链，递归会爆栈，这里统一用迭代写法。
st_u = [1]
st_edge = [head[1]]
st_sub = [1]
st_same = [1]
st_positive = [0]
st_only = [0]
st_before = [0]
parent[1] = 0

while st_u:
    u = st_u[-1]
    e = st_edge[-1]
    if e != -1:
        st_edge[-1] = nxt[e]
        v = to[e]
        if v == parent[u]:
            continue
        parent[v] = u
        st_before[-1] = done_cnt[color[u]]  # 进入儿子前的快照
        st_u.append(v)
        st_edge.append(head[v])
        st_sub.append(1)
        st_same.append(1)
        st_positive.append(0)
        st_only.append(0)
        st_before.append(0)
        continue

    # 儿子都处理完了，结算 u。
    inside = st_same.pop()
    positive = st_positive.pop()
    only_child = st_only.pop()
    sub = st_sub.pop()
    st_before.pop()  # 本帧的 before 已用完, 与 C++ 版同步出栈
    st_edge.pop()
    st_u.pop()

    c = color[u]
    subtree[u] = sub
    done_cnt[c] += 1

    if total_cnt[c] > 1:
        # 同色点分布在几个方向：含同色点的儿子各算一个，子树外还有同色点再算一个。
        outside = total_cnt[c] - inside
        directions = positive + (1 if outside > 0 else 0)
        if directions >= 3:
            too_many[c] = 1          # 三点散开，不存在装得下它们的路径
        elif directions == 1:
            # u 是最外层的端点，外侧连通块 = 含同色点那一侧的补集
            if positive == 1:
                side = n - subtree[only_child]
            else:
                side = sub
            leaf_cnt[c] += 1
            if leaf_cnt[c] == 1:
                side_first[c] = side
            elif leaf_cnt[c] == 2:
                side_second[c] = side

    if st_u:
        # 把 u 这棵子树的结果累加到父亲身上。
        p = st_u[-1]
        p_color = color[p]
        added = done_cnt[p_color] - st_before[-1]
        st_sub[-1] += sub
        st_same[-1] += added
        if added > 0:
            st_positive[-1] += 1
            st_only[-1] = u

all_pairs = n * (n - 1) // 2
for c in range(1, n + 1):
    answer[c] = all_pairs  # 没有该颜色节点时所有点对都合法

for c in range(1, n + 1):
    if total_cnt[c] == 0:
        continue
    if total_cnt[c] == 1:
        # 只有 1 个同色点 x：用补集减去「两端点在同一连通块」的点对
        x = first_node[c]
        out = n - subtree[x]
        excluded = out * (out - 1) // 2
        e = head[x]
        while e != -1:
            v = to[e]
            if parent[v] == x:
                size = subtree[v]
                excluded += size * (size - 1) // 2
            e = nxt[e]
        answer[c] = all_pairs - excluded
    else:
        # 恰好两个端点才共线，否则答案为 0
        if too_many[c] or leaf_cnt[c] != 2:
            answer[c] = 0
        else:
            answer[c] = side_first[c] * side_second[c]

out_lines = []
for c in range(1, n + 1):
    out_lines.append(str(answer[c]))
sys.stdout.write("\n".join(out_lines) + "\n")

复杂度

时间 $O(n)$ : 建树 $O(n)$ , 迭代 DFS 中每个节点进入、离开各处理常数次, 每条边被看两遍, 最后一遍输出 $O(n)$ . 没有 $\log$ 因子.

空间 $O(n)$ : 迭代栈最多 $n$ 帧, 每帧 6 个 int 加 1 个 long long(共 32 字节), 约 32MB; 邻接表 vector<int> g[n] 有 $n$ 个空的 vector 对象, 每个约 24 字节, 合计约 24MB, 加 $2(n-1)$ 条边的整数本身约 8MB; 按颜色的计数数组与 color/parent/subtree 等约 $8n$ 个 int, 约 32MB; 三个 long long 数组约 24MB. 合计约 120MB, 与实测峰值 148MB 的差距来自逐节点分配带来的容器与对齐开销.

实测 $n = 10^6$ (同一判题环境): 双色链 434ms, 单色链 252ms, 菊花 370ms, 随机树 898ms, 完全二叉树 887ms, 峰值内存 102–148MB. 相对解法二最慢的 1663ms(随机树), 时间余量从约 1.3s 变成约 2.1s; 峰值内存也从 208–219MB 降到 102–148MB.

Python 版用 array 压内存、一次性读入全部输入, 常数比 C++ 大。按洛谷对 Python 放宽 10 倍时限的口径(3s 题按 30s 计), 它在本机判题环境下跑满 $n = 10^6$ 的五种结构全部通过:

结构	C++	Python(30s 口径)
双色链	434ms	5157ms / 399MB
单色链	252ms	6962ms / 400MB
菊花	370ms	6287ms / 277MB
随机树	898ms	6420ms / 334MB
完全二叉树	887ms	7479ms / 334MB

最慢一组 7479ms 约为 30s 上限的四分之一, 内存 400MB 也在 500MB 之内, 所以这一版 Python 实现本身可以作为满分提交。如果目标判题站不对 Python 放宽时限, 同一份代码在 3s 下只覆盖 $n \lesssim 2.5\times10^5$ , 那时请改用 C++ 版。

复杂度对比

解法	时间	空间	$n=10^6$ 最慢时间(用例)	$n=10^6$ 峰值内存(用例)
一、暴力枚举	$O(n^3)$	$O(n)$	不可用	不可用
二、LCA + 直径	$O(n\log n)$	$O(n\log n)$	1663ms(随机树)	219MB(双色链)
三、数方向	$O(n)$	$O(n)$	887ms(完全二叉树)	148MB(双色链)

解法二和解法三都能通过全部数据, 差别在常数和实现负担: 解法二需要倍增表和三种 LCA 操作, 解法三只有一个遍历和几个计数器. 如果已经写过 LCA 模板, 解法二可以很快写出; 如果从头写, 解法三的代码更短, 也少一层需要验证的组件.

总结

“路径必须覆盖一组点"这一类问题, 可以固定成两步: 先把这组点压成两个端点(直径端点, 或者数方向得到的两个最外层点), 再把"路径覆盖这组点"翻译成"两端分居两个外侧连通块”, 答案就是两个连通块大小的乘积. 中间的分叉情形单独判死.

解法二用距离等式刻画端点, 依赖 LCA; 解法三用方向数刻画端点, 一次自底向上扫描就够. 两者选择的工具不同, 但最后落到同一个乘积式, 这也是本题的核心结论.

图示解析

下面的路线图对比两条正解路线从同一个瓶颈出发的不同走向.

flowchart TB
  BR["暴力枚举 brute.cpp<br/>枚举点对并收集路径节点<br/>O(n^3), 只作对拍基线"]
  BR -.->|"对拍基线, 不参与下面的分类"| CLS["共同观察: 答案只由同色点集合的形状决定<br/>按 |S_c| 分类<br/>0 个: C(n,2)<br/>1 个: 补集<br/>多个: 看是否共线"]
  CLS --> P1["解法二: 距离刻画端点<br/>两次最远点扫描得直径端点<br/>共线: dist(a,x)+dist(x,b)=dist(a,b)<br/>计数: 切断两端第一条边后相乘<br/>O(n log n) / O(n log n)"]
  CLS --> P2["解法三: 方向数刻画端点<br/>自底向上算每点的 dirs<br/>dirs>=3 判死, dirs==1 是端点<br/>端点恰好 2 个则两侧相乘<br/>O(n) / O(n)"]
  P1 --> R1["main2.cpp"]
  P2 --> R2["main.cpp (正式主解)"]

两条路线共用最上面的分类, 都终止于"两端外侧连通块大小相乘". 左侧靠距离查询把几何条件写死, 代价是每步一次 LCA; 右侧靠方向计数绕开了距离, 因此少掉 $\log$ 因子. 图中分类层是两种解法都不需要重复推导的部分, 从分类往下才是各自的选择.

目录