Sweetlemon 的博客

Sweetlemon 的博客

需要变换思路和使用技巧的点分治题——树上游戏

posted on 2019-04-12 10:46:49 | under 题解 |

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 / $919\mathrm{ms}$ (O2

带注释的代码

#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);
}