洛谷 P2664 树上游戏 题解

这道题有两种比较好的做法。除了本文介绍的做法以外，还有一种点分治的做法，在这里（不过现在还没写）。

分析

题目给出一棵有 n 个节点的树，每个节点都有一种颜色。

记节点 i 到节点 j 的路径上出现的不同颜色数量为：

$$s(i,j)$$

题目要求对每个节点 i 求出：

$$sum_i=\sum_{j=1}^{n}s(i,j)$$

如果直接枚举起点和终点，就需要处理平方级别数量的路径；而且对于每条路径，还要统计路径上出现了多少种不同颜色，显然无法通过本题。

遇到这种统计“路径上不同种类数量”的问题，可以尝试拆分贡献：不再逐条路径统计出现了多少种颜色，而是分别考虑每一种颜色对答案产生了多少贡献。

按颜色拆分贡献

先固定一个起点 i，再考虑一种颜色 c。

对于终点 j，只要从 i 到 j 的路径上出现过颜色 c，颜色 c 就会对 s(i,j) 贡献 1。

因此，颜色 c 对 sum_i 的贡献，就是满足下面条件的终点 j 的数量：

从 i 到 j 的路径上至少出现过一次颜色 c。

直接统计路径是否经过颜色 c 仍然不方便，可以考虑它的反面：

哪些终点 j 到 i 的路径完全不经过颜色 c？

将树上所有颜色为 c 的节点及其相连的边删除，原树会被分成若干个连通块。

假设节点 i 所在连通块的大小为 sz，那么：

• 连通块内的 sz 个节点到 i 的路径不经过颜色 c；
• 连通块外的节点到 i 的路径一定经过颜色 c。

所以，颜色 c 对这个连通块中每个节点的贡献都是：

$$n-sz$$

于是，问题转化为：

对每一种颜色，删除所有这种颜色的节点，求出剩余连通块的大小，并给同一连通块中的所有节点统一计算贡献。

连通块是怎样形成的

将树以节点 1 为根。

固定一种颜色 c，考虑一个颜色为 c 的节点 u。

删除 u 后，从它的某个儿子 v 向下看，v 的子树中会形成一个连通块。

如果 v 的子树中不存在其他颜色为 c 的节点，那么这个连通块就是完整的 v 子树，大小为：

$$sz_v$$

但如果 v 的子树中还存在颜色为 c 的后代节点，那么这些节点同样会被删除，并继续把 v 的子树切开。

因此，与 u 相邻并从 v 方向延伸的连通块，可以看作：

v 的整棵子树，减去其中被其他颜色为 c 的节点隔开的部分。

如下图所示：

现在的问题是，怎样在一次 DFS 中同时求出这些连通块的大小。

用全局桶维护颜色信息

如果枚举每一种颜色，并为每种颜色单独遍历整棵树，复杂度最坏会达到：

$$O(n^2)$$

为了避免枚举颜色，我们希望在遍历每个节点时，顺便处理与当前节点颜色有关的连通块。

代码中维护了一个全局桶：

colsz[c]

它记录颜色 c 在当前 DFS 过程中已经累计的信息。

考虑当前节点 u 的一个儿子 v，现在需要处理的颜色为 a[u]。

进入 v 的子树前，先保存当前的桶值：

int ts=colsz[a[u]];

然后递归处理 v：

dfs1(v,u);

从 v 的子树回溯以后，colsz[a[u]] 中会增加在 v 子树内处理颜色 a[u] 时产生的信息。

递归前后的差值为：

$$colsz[a[u]]-ts$$

这部分就是 v 子树中需要从完整子树里扣除的大小。

因此，从 v 方向得到的连通块大小为：

$$ns=sz[v]-\left(colsz[a[u]]-ts\right)$$

整理可得：

$$ns=sz[v]-colsz[a[u]]+ts$$

对应代码为：

int ts=colsz[a[u]];
dfs1(v,u);
int ns=sz[v]-colsz[a[u]]+ts;

通过记录递归前后的桶值变化，就可以在同一次 DFS 中处理所有颜色，而不需要依次枚举每种颜色再遍历整棵树。

对连通块整体修改

求出连通块大小 ns 后，还需要把它对答案产生的影响加到连通块中的所有节点上。

如果直接遍历连通块中的每个节点，复杂度仍然可能退化为平方级别。

由于树上的一棵子树对应一段完整的范围，可以使用树上差分完成整体修改。

首先执行：

cf[v]+=ns;

这表示对 v 的整棵子树增加 ns。

但 v 的子树中可能存在颜色为 a[u] 的节点。这些节点会被删除，并把当前连通块截断，所以它们下面的部分不应该获得这次增加的 ns。

因此，还需要找到 v 子树中颜色为 a[u] 的节点，并在这些节点处减去 ns。

用 DFS 序寻找同色节点

对于每一种颜色，使用一个 vector 记录这种颜色对应的节点：

vector<int> ve[N];

处理完节点 u 时，将它放入对应颜色的容器：

ve[a[u]].push_back(u);

为了判断其中哪些节点位于当前子树内，在第一次 DFS 时记录每个节点的 DFS 序：

dfn[u]=++cnt;

在深度优先遍历中，一棵子树中的节点会被连续访问。

处理完儿子 v 并回溯到 u 时，当前加入 ve[a[u]] 尾部、且 DFS 序大于 dfn[u] 的节点，就是刚才在 u 的子树中遇到的颜色为 a[u] 的后代节点。

因此可以执行：

while(ve[a[u]].size()&&dfn[ve[a[u]].back()]>dfn[u]){
    cf[ve[a[u]].back()]-=ns;
    ve[a[u]].pop_back();
}

在这些同色节点处减去 ns，可以抵消之前对子树的整体增加，使 ns 只作用于当前连通块。

处理完当前儿子对应的连通块后，将它的大小加入当前颜色的桶：

colsz[a[u]]+=ns;

所有儿子都处理完以后，当前节点 u 本身也是一个颜色为 a[u] 的节点，所以还要执行：

colsz[a[u]]++;
ve[a[u]].push_back(u);

处理根所在的连通块

前面的过程处理了从某个同色节点向其儿子方向延伸的连通块。

但是，删除一种颜色后，还可能存在一个靠近根节点的连通块。这个连通块不位于任何同色节点的儿子方向，因此不会在前面的过程中被处理。

第一次 DFS 结束后，对于颜色 c，colsz[c] 已经记录了其他部分的总大小，所以根所在连通块的大小为：

$$n-colsz[c]$$

因此，需要对每一种实际出现过的颜色额外处理一次。

首先在根节点处加上这个连通块的大小：

cf[1]+=n-colsz[i];

这样相当于先给整棵树增加：

$$n-colsz[i]$$

但颜色为 i 的节点会将树截断，这些节点下面的部分不属于根所在的连通块，所以还需要在对应节点处减去这次增加的值：

for(int j=0;j<ve[i].size();j++){
    cf[ve[i][j]]-=n-colsz[i];
}

至此，删除每一种颜色后产生的所有连通块都得到了处理。

还原树上差分

所有修改都保存在差分数组 cf 中。

第二次 DFS 从根开始向下累加：

f[u]=f[fa]+cf[u];

这样就能把树上差分还原为每个节点的实际值。

在处理完所有颜色后，f[u] 表示对于每一种颜色，节点 u 所在的不含这种颜色的连通块大小之和。

设树中实际出现的颜色数量为 num。

对于每一种颜色，如果暂时认为从 u 到所有 n 个节点的路径都包含这种颜色，那么所有颜色产生的总贡献为：

$$n\times num$$

对于颜色 c，与 u 位于同一个剩余连通块中的节点，其到 u 的路径不经过颜色 c。这些多计算的部分已经累计在 f[u] 中。

因此，节点 u 的最终答案为：

$$sum_u=n\times num-f[u]$$

对应代码为：

cout<<(ll)n*num-f[i]<<"\n";

算法流程

1. 将树以节点 1 为根。
2. 进行第一次 DFS，计算每个节点的子树大小和 DFS 序。
3. 在递归和回溯过程中，通过 colsz 的变化计算各个连通块的大小。
4. 使用 cf 对连通块进行树上差分。
5. 第一次 DFS 结束后，额外处理每种颜色对应的根部连通块。
6. 进行第二次 DFS，从根向下累加差分值。
7. 对每个节点输出：

$$n\times num-f[i]$$

复杂度分析

两次 DFS 都只遍历每个节点和每条边常数次。

对于所有颜色，所有 ve 中存储的节点总数为 n；每个节点只会被加入和弹出常数次。

因此，总时间复杂度为：

$$O(n)$$

空间复杂度为：

$$O(n)$$

点击展开代码

代码：

#include<iostream>
#include<algorithm>
#include<string.h>
#include<vector>
using namespace std;
typedef long long ll;

const int N=1e5+5;
int n,num,m,cnt;
int a[N];
bool vis[N];
int e[N<<1],ne[N<<1],h[N],idx;
int colsz[N],sz[N],dfn[N];
ll cf[N],f[N];
vector<int> ve[N];

void add(int a,int b)
{
    e[idx]=b,ne[idx]=h[a],h[a]=idx++;
}

void dfs1(int u,int fa)
{
    sz[u]=1;
    dfn[u]=++cnt;
    for(int i=h[u];~i;i=ne[i]){
        int v=e[i];
        if(v==fa) continue;
        int ts=colsz[a[u]];
        dfs1(v,u);
        int ns=sz[v]-colsz[a[u]]+ts;
        cf[v]+=ns;
        colsz[a[u]]+=ns;
        while(ve[a[u]].size()&&dfn[ve[a[u]].back()]>dfn[u]){
            cf[ve[a[u]].back()]-=ns;
            ve[a[u]].pop_back();
        }
        sz[u]+=sz[v];
    }
    colsz[a[u]]++;
    ve[a[u]].push_back(u);
}

void dfs2(int u,int fa)
{
    f[u]=f[fa]+cf[u];
    for(int i=h[u];~i;i=ne[i]){
        int v=e[i];
        if(v!=fa) dfs2(v,u);
    }
}

int main()
{
    ios::sync_with_stdio(0);
    cin.tie(0),cout.tie(0);
    memset(h,-1,sizeof h);
    cin>>n;
    for(int i=1;i<=n;i++){
        cin>>a[i];
        m=max(m,a[i]);
        vis[a[i]]=1;
    }
    for(int i=1,x,y;i<n;i++){
        cin>>x>>y;
        add(x,y),add(y,x);
    }
    dfs1(1,0);
    for(int i=1;i<=m;i++){
        if(vis[i]){
            num++;
            cf[1]+=n-colsz[i];
            for(int j=0;j<ve[i].size();j++){
                cf[ve[i][j]]-=n-colsz[i];
            }
        }
    }
    dfs2(1,0);
    for(int i=1;i<=n;i++) cout<<(ll)n*num-f[i]<<"\n";

    return 0;
}