P2664 树上游戏

题目描述

\(\text{lrb}\)有一棵树，树的每个节点有个颜色。给一个长度为\(n\)的颜色序列，定义\(s(i,j)\) 为 \(i\) 到 \(j\) 的颜色数量。以及\(sum_i=\sum\limits_{j=1}^ns(i,j)\)

现在他想让你求出所有的\(sum_i\)

输入输出格式

输入格式：

第一行为一个整数\(n\)，表示树节点的数量

第二行为\(n\)个整数，分别表示\(n\)个节点的颜色\(c[1],c[2],\dots,c[n]\)

接下来\(n-1\)行，每行为两个整数\(x,y\)，表示\(x\)和\(y\)之间有一条边

输出格式：

输出\(n\)行，第\(i\)行为\(sum[i]\)

说明

对于\(40\%\)的数据，\(n\le 2000\)

对于\(100\%\)的数据，\(1\le n,c[i]\le10^5\)

---

本辣鸡深深的认识到了自己的淀粉质写法有多么的诡异..

这个题在淀粉质上的思路还是比较简单的，就是写起来很烦人。

我的做法大致是

每次出去遍历子树时维护一个当前链颜色集合和一个过点分治根的每个颜色所属链的个数。然后每个点根据个数加一些本身颜色的贡献就可以了。注意要从左从右各做一遍。

说起来比较简单，但实际上还是要想一想的。

然后我拿\(map\)乐呵呵的维护颜色集合然后成功T飞

发现拿数组就可以了，但是最后要dfs去清空数组。

---

Code:

#include <cstdio>

#include <vector>

#include <map>

#include <cctype>

#define ll long long

const int N=1e5+10;

int read()

{

    int x=0;char c=getchar();

    while(!isdigit(c)) c=getchar();

    while(isdigit(c)) {x=x*10+c-'0';c=getchar();}

    return x;

}

std::vector <int> Edge[N];

ll sum[N];

const int inf=0x3f3f3f3f;

int n,c[N],siz[N],del[N],mi,rt,lassiz,typ,sumnow;

void getroot(int now,int fa,int s)

{

    int mx=0;

    siz[now]=1;

    for(int i=0;i<Edge[now].size();i++)

    {

        int v=Edge[now][i];

        if(del[v]||v==fa) continue;

        getroot(v,now,s);

        siz[now]+=siz[v];

        mx=mx>siz[v]?mx:siz[v];

    }

    mx=mx>s-siz[now]?mx:s-siz[now];

    if(mi>mx) rt=now,mi=mx;

}

int nowcolor[N];//存到根的颜色集合

ll lascolor[N],tlascolor[N];//存外面的链每种颜色在多少条链出现过

void dfs(int now,int fa,int totcolor,ll outsum)

{

    siz[now]=1;int flag=0;

    if(!nowcolor[c[now]])//如果到根第一次出现这个颜色

    {

        ++totcolor,nowcolor[c[now]]=1;

        outsum=outsum+lassiz-lascolor[c[now]];//有这么多条路径没有这种颜色

        flag=1;

    }

    sum[now]+=outsum;

    sum[now]=sum[now]-totcolor*typ;//减去统计自己与根的答案

    sumnow+=totcolor;

    for(int i=0;i<Edge[now].size();i++)

    {

        int v=Edge[now][i];

        if(del[v]||v==fa) continue;

        dfs(v,now,totcolor,outsum);

        siz[now]+=siz[v];

    }

    if(flag)//如果到根第一次出现这个颜色

    {

        nowcolor[c[now]]=0;

        tlascolor[c[now]]+=1ll*siz[now];

    }

}

void dfsin(int now,int fa)

{

    int flag=0;

    if(!nowcolor[c[now]])//如果到根第一次出现这个颜色

    {

        nowcolor[c[now]]=1;

        flag=1;

    }

    for(int i=0;i<Edge[now].size();i++)

    {

        int v=Edge[now][i];

        if(del[v]||v==fa) continue;

        dfsin(v,now);

    }

    if(flag)//如果到根第一次出现这个颜色

    {

        nowcolor[c[now]]=0;

        lascolor[c[now]]+=1ll*siz[now];

    }

}

void dfsclear(int now,int fa)

{

    lascolor[c[now]]=0;

    for(int i=0;i<Edge[now].size();i++)

    {

        int v=Edge[now][i];

        if(del[v]||v==fa) continue;

        dfsclear(v,now);

    }

}

void divide(int now,int s)

{

    mi=inf;getroot(now,now,s);

    del[now=rt]=1;

    lascolor[c[now]]=1,typ=0,sumnow=1,lassiz=1,nowcolor[c[now]]=1;

    for(int i=0;i<Edge[now].size();i++)

    {

        int v=Edge[now][i];

        if(del[v]) continue;

        dfs(v,now,1,sumnow);

        lascolor[c[now]]+=1ll*siz[v];

        dfsin(v,now);

        lassiz+=siz[v];

    }

    sum[now]=sum[now]+sumnow;

    dfsclear(now,now);

    lascolor[c[now]]=1,typ=1,sumnow=1,lassiz=1;

    for(int i=Edge[now].size()-1;~i;i--)

    {

        int v=Edge[now][i];

        if(del[v]) continue;

        dfs(v,now,1,sumnow);

        lascolor[c[now]]+=1ll*siz[v];

        dfsin(v,now);

        lassiz+=siz[v];

    }

    nowcolor[c[now]]=0;

    dfsclear(now,now);

    for(int i=0;i<Edge[now].size();i++)

    {

        int v=Edge[now][i];

        if(!del[v]) divide(v,siz[v]);

    }

}

int main()

{

    n=read();

    for(int i=1;i<=n;i++) c[i]=read();

    for(int u,v,i=1;i<n;i++)

    {

        u=read(),v=read();

        Edge[u].push_back(v);

        Edge[v].push_back(u);

    }

    divide(1,n);

    for(int i=1;i<=n;i++) printf("%lld\n",sum[i]);

    return 0;

}

---

2018.12.4