[JZOJ5058]采蘑菇

最新推荐文章于 2022-10-30 18:00:00 发布

原创最新推荐文章于 2022-10-30 18:00:00 发布 · 948 阅读

0 ·

CC 4.0 BY-SA版权

文章标签：

#OI #点分治

纪中OJ 同时被 2 个专栏收录

127 篇文章

订阅专栏

点分治

11 篇文章

订阅专栏

本文介绍了一种基于点分治的算法，用于解决特定类型的树形结构问题，具体是在一棵带有节点颜色信息的树上进行路径颜色种类数的统计。

题目大意

给定一棵 $n$ 个节点的树，每个点有一个颜色种类 $c_i$ 。
对于每一个点 $x$ ，你需要统计从 $x$ 出发的所有路径的颜色种类数之和。

$1\le n\le3\times10^5,0\le c_i\le n$

题目分治

首先这题虚树肯定可以做，这里不讲。
考虑使用点分治，先不考虑有多种颜色。假设我只想统计出现过某一种颜色的路径总数。
对于分治重心 $c$ ，在分治过程中做到点 $x$ ：
$\bullet\$ 如果 $x$ 到 $c$ 的路径上已经有了这一种颜色，那么 $x$ 的答案显然就要加上当前分治层的点数减去 $x$ 所在子树的点数。
$\bullet\$ 如果 $x$ 到 $c$ 的路径上没有这种颜色，我们就要通过预处理求出所有到 $c$ 路径（不包括 $c$ ）包含该种颜色的路径条数统计出来，减去 $x$ 所在子树中满足同样条件的路径数，加在 $x$ 的答案上面。
现在考虑将其扩展到多种颜色上。
对于分治重心 $c$ ，在分治过程中做到点 $x$ ：
$\bullet\$ 设其到 $c$ 路径上颜色种类数为 $cnt$ ，分治层点数减去 $x$ 所在子树的点数是 $siz$ ，答案就要加上 $cnt\times siz$ 。
$\bullet\$ 对于那些没有出现过的颜色种类，我们考虑先预处理不包含颜色 $i$ 的到 $c$ 路径（不包含 $c$ ）条数 $f_i$ ，那么答案就要加上每种没有出现过的颜色的 $f$ 值，算的时候由于我们是深搜实现，一开始用 $sum$ 记录 $f$ 值和，每新出现一种颜色就减掉对应的 $f$ ，消失一种颜色就加上，这样就可以快速求出要加的值了。当然，我们还要减去和 $x$ 同一棵子树的路径，这个在实现的时候，我们再深搜计算一个子树之前先深搜一边把这个子树的路径从 $f$ 中删掉就好了，深搜计算完之后再深搜加回去就好了。
时间复杂度 $\mathrm O(n\log n)$ 。

代码实现

#include <iostream>
#include <cstdio>
#include <cctype>

using namespace std;

typedef long long LL;

int read()
{
    int x=0,f=1;
    char ch=getchar();
    while (!isdigit(ch)) f=ch=='-'?-1:f,ch=getchar();
    while (isdigit(ch)) x=x*10+ch-'0',ch=getchar();
    return x*f;
}

int buf[30];

void write(LL x)
{
    if (x<0) putchar('-'),x=-x;
    for (;x;x/=10) buf[++buf[0]]=x%10;
    if (!buf[0]) buf[++buf[0]]=0;
    for (;buf[0];putchar('0'+buf[buf[0]--]));
}

const int N=300050;
const int E=N<<1;

int last[N],fa[N],size[N],col[N],f[N],ext[N],que[N];
int nxt[E],tov[E];
bool vis[N];
LL ans[N];
int n,tot,head,tail,cur;
LL sum;

void insert(int x,int y){tov[++tot]=y,nxt[tot]=last[x],last[x]=tot;}

int core(int og)
{
    int i,x,y,ret,rets=n,tmp;
    for (head=0,fa[que[tail=1]=og]=0;head<tail;)
        for (size[x=que[++head]]=1,i=last[x];i;i=nxt[i])
            if ((y=tov[i])!=fa[x]&&!vis[y])
                fa[que[++tail]=y]=x;
    for (head=tail;head>1;--head) size[fa[que[head]]]+=size[que[head]];
    for (head=1;head<=tail;++head)
    {
        for (tmp=size[og]-size[x=que[head]],i=last[x];i;i=nxt[i])
            if ((y=tov[i])!=fa[x]&&!vis[y]) tmp=max(tmp,size[y]);
        if (tmp<rets) ret=x,rets=tmp;
    }
    return ret;
}

void dfs(int x)
{
    size[x]=1;
    for (int i=last[x],y;i;i=nxt[i])
        if ((y=tov[i])!=fa[x]&&!vis[y])
            fa[y]=x,dfs(y),size[x]+=size[y];
}

void count(int x,int *f,int sig)
{
    if (!ext[col[x]]++) f[col[x]]+=sig*size[x],sum+=sig*size[x],++cur;
    for (int i=last[x],y;i;i=nxt[i])
        if ((y=tov[i])!=fa[x]&&!vis[y]) count(y,f,sig);
    if (!--ext[col[x]]) --cur;
}

void calc(int x,int siz,int c)
{
    if (!ext[col[x]]++) sum-=f[col[x]],++cur;
    ans[x]+=1ll*siz*cur+sum,ans[c]+=cur;
    for (int i=last[x],y;i;i=nxt[i])
        if ((y=tov[i])!=fa[x]&&!vis[y]) calc(y,siz,c);
    if (!--ext[col[x]]) sum+=f[col[x]],--cur;
}

void solve(int x)
{
    int c=core(x);
    ++ans[c],size[c]=1;
    for (int i=last[c],y;i;i=nxt[i])
        if (!vis[y=tov[i]]) fa[y]=c,dfs(y),size[c]+=size[y];
    for (int i=last[c],y;i;i=nxt[i])
        if (!vis[y=tov[i]]) count(y,f,1);
    for (int i=last[c],y;i;i=nxt[i])
        if (!vis[y=tov[i]]) count(y,f,-1),++ext[col[c]],sum-=f[col[c]],cur=1,calc(y,size[c]-size[y],c),--ext[col[c]],sum+=f[col[c]],cur=0,count(y,f,1);
    for (int i=last[c],y;i;i=nxt[i])
        if (!vis[y=tov[i]]) count(y,f,-1);
    vis[c]=1;
    for (int i=last[c],y;i;i=nxt[i])
        if (!vis[y=tov[i]]) solve(y);
}

int main()
{
    freopen("mushroom.in","r",stdin),freopen("mushroom.out","w",stdout);
    n=read();
    for (int i=1;i<=n;++i) col[i]=read();
    for (int i=1,x,y;i<n;++i) x=read(),y=read(),insert(x,y),insert(y,x);
    solve(1);
    for (int i=1;i<=n;++i) write(ans[i]),putchar('\n');
    fclose(stdin),fclose(stdout);
    return 0;
}