Luogu P5588 小猪佩奇爬树 题解
写在前面
我发现这个题目题解是少之又少……于是你就看到了本篇题解。
笔者写下这篇文章，一是为了总结，二是希望帮助到更多的人。
希望大家看完这篇文章都能有所收获。
题目传送门 P5588 小猪佩奇爬树
题解部分
最初的思路
一开始我看到本题，就想是不是和 \(LCA\) 有关。
当时我思考只要求出每个颜色的 \(LCA\) 即伸展的最下方节点（包含该颜色的节点）
只要所有颜色在一条链上，我们就可以统计答案。特别的，一个点也算链。
但是实际上并不需要这么复杂，所以我月赛也没有打出本题。
思路改进及代码实现
那么我们先考虑简单的情况，就是没有这个颜色的节点，那么我们答案即为 \(\frac{n*(n-1)}{2}\)。
还有就是这个颜色只有一个节点，那么就是统计一个这节点为根的所有子树大小互乘累加。
这个是根据乘法原理得来的，如果你没有理解上面那句话，那么你可以看下面这句代码来理解：
void dfs(int x, int fa) {
……;
size[x] = 1;
for(int i = head[x]; i; i = Next[i]) {
if(ver[i] == fa) continue;
dfs(y, x);
ans1[x] += 1LL * size[x] * size[y];
size[x] += size[y];
……;
} ans1[x] += size[x] * (n - size[x]);
……;
}其中 \(ans1\)即为我们所要的答案（\(size\)统计子树大小）。
我们只需要再开一个额外的数组确认当前颜色的 \(x\) 在哪即可。
那么以上就是一个颜色只有一个节点的情况。
下面重点讨论一个颜色有两个以上节点的情况，那么可能会出现两种情况：
所有这样的节点都在一条链上，也就是说这个颜色有两个端点。
所有这样的颜色不在一条链上，也就是说有三个及以上端点。
那么显而易见的，对于第2种情况一定不存在一种合法的解，输出 \(0\) 即可。
所以重点是解决第1种情况。
我们考虑对于一个颜色，什么时候可能会是端点。我们记一个数组 \(cnt\) 表示颜色 \(c\) 到目前为止出现的次数。
那么对于树上操作，一般进行 \(dfs\) ，那么 \(cnt[c]\)(\(c\) 是当前颜色)就极可能直接发生改变。
那么对于每一次改变，我们定义一个变量 \(flag++\)，只要最终结果 \(flag=1\) 就极可能是个端点。
还有另一种可能，就是进入当前节点时 \(cnt[c]\) 已经有值，或者当前节点不为当前颜色的最后一个节点，那么也要使 \(flag++\)。
具体再解释下 \(flag\) 可能会更好理解：\(flag\) 代表当前节点下子树颜色也为 \(c\) 的个数。
这样好理解了吧？记得最后处理完上面的步骤还要让 \(cnt[c]++\)一次（当前颜色也是 \(c\) 啊 \(QAQ\)）。
只要最后 \(flag=1\)那么即为端点。第一次进入直接让当前颜色的节点指针赋给 \(nos\) (记录当前颜色从哪开始)
然后答案即为最两端的节点的子树大小之积。
- 很抽象？看看怎么实现的吧：
int color[N], tot[N], cnt[N], n;
int enos[N], size[N], nos[N];
long long ans1[N], ans2[N];
// ans1 统计一个颜色一个节点的情况
// ans2 统计一个颜色多个节点的情况
// enos[c] 表示 c 的个数 nos[c] 表示 c 对应的位置
// 至于 nos 的用处过会儿完整代码你就明白
// tot[c] 表示颜色 c 总共的个数 剩下的含义看上面啦
inline void dfs(int x, int fa) {
int c = color[x], k = cnt[c];
int flag = 0, pos = 0;
size[x] = 1;
for(int i = head[x]; i; i = Next[i]) {
int y = ver[i];
if(y == fa) continue;
int lastans = cnt[c];
dfs(y, x);
ans1[x] += 1LL * size[x] * size[y];
size[x] += size[y];
if(lastans != cnt[c]) flag++, pos = y;
}
ans1[x] += 1LL * size[x] * (n - size[x]);
if(k || cnt[c] != tot[c] - 1) flag++;
cnt[c]++; // 当前节点颜色是 c
if(flag == 1) { // 端点
if(!enos[c]) nos[c] = x;
else {
int p = pos ? n - size[pos] : size[x];
ans2[c] = 1LL * size[nos[c]] * p;
} enos[c]++;
}
}这样子大家应该可以理解如何解决本题了吧？
那么分析下这个 \(dfs\) 的复杂度，很明显是 \(O(N)\) 的……（为什么某位神仙打了\(O(N \log N)\)的……）
好了下面放上完整代码吧（建议按照思路自己先打，不知道 \(nos\) 怎么用可以看看）
\(Code:\)
#include<bits/stdc++.h>
using namespace std;
const int N = 1e6 + 10;
int head[N], Next[N << 1], ver[N << 1];
void add(int x, int y) {
static int cnt = 0;
ver[++cnt] = y, Next[cnt] = head[x], head[x] = cnt;
ver[++cnt] = x, Next[cnt] = head[y], head[y] = cnt;
}
int color[N], tot[N], cnt[N], n;
int enos[N], size[N], nos[N];
long long ans1[N], ans2[N];
inline void dfs(int x, int fa) {
int c = color[x], k = cnt[c];
int flag = 0, pos = 0;
size[x] = 1;
for(int i = head[x]; i; i = Next[i]) {
int y = ver[i];
if(y == fa) continue;
int lastans = cnt[c];
dfs(y, x);
ans1[x] += 1LL * size[x] * size[y];
size[x] += size[y];
if(lastans != cnt[c]) flag++, pos = y;
}
ans1[x] += 1LL * size[x] * (n - size[x]);
if(k || cnt[c] != tot[c] - 1) flag++;
cnt[c]++;
if(flag == 1) { // 端点
if(!enos[c]) nos[c] = x;
else {
int p = pos ? n - size[pos] : size[x];
ans2[c] = 1LL * size[nos[c]] * p;
} enos[c]++;
}
}
int main() {
scanf("%d", &n);
for(int i = 1; i <= n; i++) {
scanf("%d", color + i);
tot[color[i]]++;
nos[color[i]] = i;
}
for(int i = 1, x, y; i < n; i++) {
scanf("%d %d", &x, &y);
add(x, y);
}
dfs(1, 0);
for(int i = 1; i <= n; i++) {
if(tot[i] == 0)
printf("%lld\n", 1LL * n * (n - 1) >> 1);
else if(tot[i] == 1)
printf("%lld\n", ans1[nos[i]]);
else if(enos[i] == 2)
printf("%lld\n", ans2[i]);
else puts("0");
}
return 0;
}写在后面
很好的一道题目，但是因为脑抽又想了很久没打出来……
所以 \(Orz\) ……