需要变换思路和使用技巧的点分治题——树上游戏
Sweetlemon
2019-04-12 10:46:49
[P2664 树上游戏](https://www.luogu.org/problemnew/show/P2664)
一道黑题，难度名副其实。
下面简述我的做法。
我的做法可能相当劣，时间复杂度是$\text{O}(n \log n)$，比不过$\text{O}(n)$的方法；而且代码比较长，达到了$210$行（带注释$242$行）；使用了很多次$\mathrm{dfs}$，常数也不够优。但是也许我的思路会对您有所启发。
看到树上“路径”，且树形$\text{dp}$不好处理，于是我们考虑使用点分治。
“颜色数”一直是难以处理的统计量，这个统计量阻止了我们直接对“节点到根的路径”进行合并，而必须考虑其他的方法。
假设我们已经分治到原树的某一部分，并且我们当前只考虑经过（分治的）根（即当前重心，不妨记为$\mathrm{root}$）的路径。那么我们需要思考两个问题：
1. 如何计算当前分治区域内根的答案？
2. 如何计算当前分治区域内其他点的答案？
为了方便处理，我们把根的颜色单独拿出来讨论。
对于根本身，根的颜色的贡献是$\mathrm{size}(\mathrm{root})$；对于根的任何一个子孙顶点，根的颜色的贡献是$\mathrm{size}(\mathrm{root})-\mathrm{size}(\mathrm{branch})$，其中$\mathrm{branch}$是指该节点所在的根的子树。
接下来，如果遇到与根同色的节点，我们就当作无色处理。
第一个问题比较容易解决。
假设我们从根走到**某一**子孙节点，在走的过程中，如果出现了某种**新**颜色，那么这条路径答案就应该$+1$。也即我们统计的思路是，只在颜色第一次出现时计算。
如果在$\mathrm{root}\rightarrow x$的路径上，$\mathrm{color}(x)$只出现一次（即$x$的祖先都与$x$不同色），那么祖先到以$x$为根的子树内的所有点的路径上都会出现$\mathrm{color}(x)$，且第一次出现的位置是$x$。从而，$x$对根的答案的贡献为$x$所在子树内顶点的个数，即$\mathrm{size}(x)$。否则如果$x$与某一祖先同色，那么$x$的颜色就不是第一次出现，对根的答案没有贡献。
不妨设节点$x$的上述贡献为$f(x)$，那么$f(x)$为$\mathrm{size}(x)$（$\mathrm{color}(x)$第一次出现）或$0$（$x$为根或$\mathrm{color}(x)$不是第一次出现)。我们把它们累加，得到整棵树的贡献和$\sum{f}$。
那么第一个问题就解决了。该区域内根的答案$\mathrm{ans}(\mathrm{root})=\mathrm{size}(\mathrm{root})+\sum{f}$，即根节点处的答案就等于自身颜色的贡献与子孙节点颜色贡献之和。
第二个问题比较困难。
点分治的合并思想启示我们，借助刚才计算的$f$解决第二个问题。我们把路径$x\rightarrow y$拆成$x\rightarrow \mathrm{root}$和$\mathrm{root} \rightarrow y$两段。
对于某一定点$x$，如果只考虑所有的$x\rightarrow \mathrm{root}$，那么这些路径一共有$\mathrm{size}(\mathrm{root})-\mathrm{size}(\mathrm{branch})$条（每个$y$对应一条），每条路径完全相同，出现的不同颜色的个数都相等；于是，路径上出现的每一种颜色$t$对$x$的答案的贡献均为$\mathrm{size}(\mathrm{root})-\mathrm{size}(\mathrm{branch})$。
如果只考虑所有的$\mathrm{root} \rightarrow y$，这些路径中每条路径出现颜色数之和即为$\sum_{y\notin \mathrm{branch}}{f(y)}=\sum{f}-\sum_{i \in \mathrm{branch}}{f(i)}$，即$\sum{f}$减去所有$x$所在分支的点的$f$值。
但是，$x\rightarrow \mathrm{root}$和$\mathrm{root} \rightarrow y$中可能会有重复的颜色。具体地，如果某种颜色$t$在$x\rightarrow \mathrm{root}$中出现了，那么对于所有的在分支外部的、$\mathrm{color}(u)=t$的顶点$u$，它们的$f$值都不能再计入到答案里，因为这一种颜色已经在$x\rightarrow \mathrm{root}$上统计过了。
如果我们记$A$为在$x\rightarrow \mathrm{root}$上出现过的颜色的集合，那么答案就是
$$\sum{f}-\sum_{i \in \mathrm{branch}}{f(i)}-\sum_{u\notin \mathrm{branch}\text{且}\mathrm{color}(u)\in A}f(u)+\sum_{t\in A}(\mathrm{size}(\mathrm{root})-\mathrm{size}(\mathrm{branch}))$$
上式已经包括了根的颜色对$x$的答案的贡献。
观察上式，我们发现要记录的量有：
1. $\sum{f}$。
2. $\mathrm{branch}$中所有点的$f$和，即$\sum_{i \in \mathrm{branch}}{f(i)}$。
3. 不在$\mathrm{branch}$中，但其颜色在$x\rightarrow \mathrm{root}$上出现过的点的$f$和，即$\sum_{u\notin \mathrm{branch}\text{且}\mathrm{color}(u)\in A}f(u)$。
4. $x\rightarrow \mathrm{root}$上所有颜色的贡献和，即$\sum_{t\in A}(\mathrm{size}(\mathrm{root})-\mathrm{size}(\mathrm{branch}))$。
对此，找到重心后，我们再进行三次$\mathrm{dfs}$。
第一次$\mathrm{dfs}$的任务是，计算出$\mathrm{size},\sum{f},\sum_{i \in \mathrm{branch}}{f(i)}$；并对每一种颜色$t$，计算所有颜色为$t$的节点的$f$和，即$\sum_{\mathrm{color}(u)=t}f(u)$。
第一次$\mathrm{dfs}$后，我们就可以得到$\mathrm{ans}(\mathrm{root})$了。
接下来开始枚举$\mathrm{root}$的子树，即枚举$\mathrm{branch}$。
对于每一个$\mathrm{branch}$，先进行第二次$\mathrm{dfs}$，任务是对每一种$\mathrm{branch}$中出现的颜色$t$，计算分支内所有颜色为$t$的节点$v$的$f$和，即$\sum_{v\in \mathrm{branch}\text{且}\mathrm{color}(v)=t}f(v)$。这样，我们就可以得到上面的第三个量：
$$\sum_{u\notin \mathrm{branch}\text{且}\mathrm{color}(u)=t}f(u)=\sum_{\mathrm{color}(u)=t}f(u)-\sum_{v\in \mathrm{branch}\text{且}\mathrm{color}(v)=t}f(v)$$
紧接着，进行第三次$\mathrm{dfs}$，任务便是计算答案。
先处理$\mathrm{ans}$式子的第二项，即我们把分支内的$f$从$\sum{f}$里暂时减掉，即把$\sum{f}$扣掉$\sum_{i \in \mathrm{branch}}{f(i)}$，待$\mathrm{dfs}$结束后再加回来。
第三、第四项合并处理。在$\mathrm{dfs}$的过程中，如果到达了一个与祖先都不同色的顶点$v$，就说明$t=\mathrm{color}(v)$是新出现的，那么$v$子树中的所有节点的集合$A$都会新增一个元素$t$。这样，第三项就增加了$\sum_{u\notin \mathrm{branch}\text{且}\mathrm{color}(u)=t}f(u)$，也即答案还需要扣除$\sum_{\mathrm{color}(u)=t}f(u)-\sum_{v\in \mathrm{branch}\text{且}\mathrm{color}(v)\in A}f(v)$；而第四项就增加了一个$\mathrm{size}(\mathrm{root})-\mathrm{size}(\mathrm{branch})$，也即答案还需要增加$\mathrm{size}(\mathrm{root})-\mathrm{size}(\mathrm{branch})$。那么我们把上述变化处理到对存储$\sum{f}$的变量处，函数即将返回时再恢复回来。
而如果$x$的颜色已经出现过了，那么我们不做上述处理。
经过这样的处理，$\mathrm{ans}$式子的每一项都在存储$\sum{f}$的变量处计算了，那么该变量的值就等于$\mathrm{ans}(x)$。
经过这样的遍历，我们就把所有答案都计算完毕了。
如上所述，我们按照点分治的基本思想，在每一层分治都作上述计算，把每一层的答案累加即可。
但在代码实现上还有一个较大的问题。由于点分治需要多次计算，因此数组必须及时清零，否则答案会不正确。
存储每个分支的$f$和的数组只需在第三次$\mathrm{dfs}$结束后顺手将该分支的位置设为$0$即可。对于存储每种颜色的$f$和的数组，可以用一个栈记录该层分治所出现的颜色，在往下递归分治前把栈中所有颜色的值清零即可。
而记录某特定分支中每一种颜色的$f$值和的数组就比较尴尬，因为它必须在计算每一个分支完毕后立即清零。我采用的方法是在第三次$\mathrm{dfs}$后再进行第四次$\mathrm{dfs}$，任务是清零该数组。
这样，这道题就完全解决了。当然，代码实现上还有许多细节，需要多多注意。例如，答案可能很大，需要开`long long`（否则只有$80$分）。
我的实现总用时[$1421\mathrm{ms}$（无`O2`）](https://www.luogu.org/record/show?rid=18133128) / [$919\mathrm{ms}$（`O2`）](https://www.luogu.org/record/show?rid=18133134)。
[带注释的代码](https://www.luogu.org/recordnew/show/18135537)
```cpp
#include <cstdio>
#include <cctype>
#define MAXN 100005
#define MAXM 200005
#define MAXIOLG 25
using namespace std;
typedef long long ll; //答案可能很大，记得开long long
//下面为目隐团部分成员/读入输出优化
typedef ll io_t;
//如月伸太郎, shintaro, No.7, 目缠, 记忆数字
io_t shin[MAXIOLG];
//濑户幸助, seto, No.2, 目盗, 读入优化
io_t seto(void);
//楯山文乃, ayano, No.0, 目挂, 输出优化
void ayano(io_t x,char spliter='\n');
int fst[MAXN],nxt[MAXM],edges=0; //存图邻接表
int g[MAXM]; //图
int sz[MAXN],mnsz,root,tree_sz; //size数组, mnsz用于找重心, root为当前分治根, tree_sz为当前分治区域点数
int visited[MAXN]; //visited记录点是否已经被删除
int color[MAXN]; //每个点的颜色
int tree_color[MAXN],tree_color_n; //栈结构, 存储当前分治区域出现的颜色, 便于清空数组
//依次为每种颜色的f和、该种颜色在之前是否出现过、该分支每种颜色的f和、每一分支的f和、存储sum(f)的变量
ll color_f[MAXN],has_col[MAXN],branch_color_f[MAXN],branch_f[MAXN],tf;
//答案数组
ll ans[MAXN];
void addedge(int u,int v); //加边
void solve(int x); //分治函数
void find_g(int x,int pa); //找重心
void dfs1(int x,int pa,int branch); //第一次dfs
void dfs2(int x,int pa,int branch); //第二次dfs
void dfs3(int x,int pa,int branch); //第三次dfs
void dfs4(int x,int pa,int branch); //第四次dfs
int main(void){
int n;
//读入
n=seto();
for (int i=1;i<=n;i++)
color[i]=seto();
for (int i=1;i<n;i++){
int u,v;
u=seto(),v=seto();
addedge(u,v);
addedge(v,u);
}
//点分治
sz[1]=n;
solve(1);
//输出
for (int i=1;i<=n;i++)
ayano(ans[i]);
return 0;
}
void solve(int x){
//分治函数
mnsz=tree_sz=sz[x];
find_g(x,0); //找重心
//清空栈
tree_color_n=0;
const int root_col=color[root];
tree_color[tree_color_n++]=root_col;
has_col[root_col]=1;
tf=0;
dfs1(root,0,0); //第一次dfs计算sz, tf, color_f, branch_f
ans[root]+=(tf+tree_sz); //得到根的答案
for (int ei=fst[root];ei;ei=nxt[ei]){
//枚举分支
int v=g[ei];
if (visited[v])
continue;
tf-=branch_f[v]; //扣除该分支f值（答案式第二项）
dfs2(v,root,v); //第二次dfs计算branch_color_f
dfs3(v,root,v); //第三次dfs计算ans
dfs4(v,root,v); //第四次dfs清零branch_color_f
tf+=branch_f[v]; //消除第二项的影响, 恢复tf变量
branch_f[v]=0; //清零branch_f
}
has_col[root_col]=0; //清零has_col
//清零color_f
while (tree_color_n){
tree_color_n--;
color_f[tree_color[tree_color_n]]=0;
}
//递归分治
visited[root]=1;
for (int ei=fst[root];ei;ei=nxt[ei]){
int v=g[ei];
if (visited[v])
continue;
solve(v);
}
}
void dfs1(int x,int pa,int branch){
//第一次dfs计算sz, tf, color_f, branch_f
if (pa==root)
branch=x; //处理branch变量
sz[x]=1; //sz初始化
int is_new_col=0,tcol=color[x];
if (!has_col[tcol]) //新出现的颜色, 先记录
is_new_col=has_col[tcol]=1;
for (int ei=fst[x];ei;ei=nxt[ei]){
int v=g[ei];
if (visited[v] || v==pa)
continue;
dfs1(v,x,branch);
sz[x]+=sz[v];
}
if (is_new_col){
//新出现的颜色, 更新变量
if (!color_f[tcol]) //颜色在整个分治区域第一次出现, 入栈
tree_color[tree_color_n++]=tcol;
color_f[tcol]+=sz[x]; //将贡献累加到color_f中
branch_f[branch]+=sz[x]; //将贡献累加到branch_f中
tf+=sz[x]; //将贡献累加到tf中
has_col[tcol]=0; //清掉has_col
}
}
void dfs2(int x,int pa,int branch){
//第二次dfs计算branch_color_f
int tcol=color[x],is_new_color=0;
if (!has_col[tcol]) //新出现的颜色, 先记录
has_col[tcol]=is_new_color=1;
for (int ei=fst[x];ei;ei=nxt[ei]){
int v=g[ei];
if (visited[v] || v==pa)
continue;
dfs2(v,x,branch);
}
if (is_new_color) //新出现的颜色, 清掉has_col并将贡献累加到branch_color_f中
has_col[tcol]=0,branch_color_f[tcol]+=sz[x];
}
void dfs3(int x,int pa,int branch){
//第三次dfs计算ans
ans[x]+=(tree_sz-sz[branch]); //根处的答案在前面未加入, 现在额外算
int tcol=color[x],is_new_color=0;
if (!has_col[tcol]){
has_col[tcol]=is_new_color=1; //新出现的颜色,记录
tf-=(color_f[tcol]-branch_color_f[tcol]); //答案式子第三项
tf+=(tree_sz-sz[branch]); //答案式子第四项
}
ans[x]+=tf; //累加答案
for (int ei=fst[x];ei;ei=nxt[ei]){
int v=g[ei];
if (visited[v] || v==pa)
continue;
dfs3(v,x,branch);
}
if (is_new_color){
//新出现的颜色,清掉has_col并恢复tf变量原来的值
has_col[tcol]=0;
tf+=(color_f[tcol]-branch_color_f[tcol]);
tf-=(tree_sz-sz[branch]);
}
}
void dfs4(int x,int pa,int branch){
//第四次dfs清零branch_color_f
branch_color_f[color[x]]=0;
for (int ei=fst[x];ei;ei=nxt[ei]){
int v=g[ei];
if (visited[v] || v==pa)
continue;
dfs4(v,x,branch);
}
}
void find_g(int x,int pa){
//找重心
sz[x]=1; //初始化size
int mxsz=0; //x的子树或x以上部分的最大size
for (int ei=fst[x];ei;ei=nxt[ei]){
int v=g[ei];
if (visited[v] || v==pa)
continue;
find_g(v,x);
sz[x]+=sz[v];
if (sz[v]>mxsz)
mxsz=sz[v];
}
if (tree_sz-sz[x]>mxsz)
mxsz=tree_sz-sz[x];
if (mxsz<mnsz)
mnsz=mxsz,root=x;
}
void addedge(int u,int v){
//加边函数
edges++;
g[edges]=v;
nxt[edges]=fst[u],fst[u]=edges;
}
//下面为目隐团部分成员/读入输出优化
io_t seto(void){
//濑户幸助, seto, No.2, 目盗, 读入优化
io_t ans=0;
int symbol=0;
char ch=getchar();
while (!isdigit(ch))
(ch=='-')?(symbol=1):(0),ch=getchar();
while (isdigit(ch))
(ans=ans*10+(ch-'0')),ch=getchar();
return (symbol)?(-ans):(ans);
}
void ayano(io_t x,char spliter){
//楯山文乃, ayano, No.0, 目挂, 输出优化
if (!x){
putchar('0'),putchar(spliter);
return;
}
if (x<0)
putchar('-'),x=-x;
int len=0;
while (x){
io_t d=x/10;
shin[len++]=x-d*10; //如月伸太郎, shintaro, No.7, 目缠, 记忆数字
x=d;
}
while (len){
len--;
putchar(shin[len]+'0');
}
putchar(spliter);
}
```