LG2664 树上游戏

树上游戏

题目描述

lrb有一棵树，树的每个节点有个颜色。给一个长度为n的颜色序列，定义s(i,j) 为i 到j 的颜色数量。以及

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

现在他想让你求出所有的sum[i]

输入输出格式

输入格式：

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

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

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

输出格式：

输出n行，第i行为sum[i]

输入输出样例

输入样例#1：
复制

5
1 2 3 2 3
1 2
2 3
2 4
1 5

输出样例#1：
复制

10
9
11
9
12

说明

sum[1]=s(1,1)+s(1,2)+s(1,3)+s(1,4)+s(1,5)=1+2+3+2+2=10
sum[2]=s(2,1)+s(2,2)+s(2,3)+s(2,4)+s(2,5)=2+1+2+1+3=9
sum[3]=s(3,1)+s(3,2)+s(3,3)+s(3,4)+s(3,5)=3+2+1+2+3=11
sum[4]=s(4,1)+s(4,2)+s(4,3)+s(4,4)+s(4,5)=2+1+2+1+3=9
sum[5]=s(5,1)+s(5,2)+s(5,3)+s(5,4)+s(5,5)=2+3+3+3+1=12

对于40%的数据，n<=2000

对于100%的数据，1<=n,c[i]<=10^5

题解

这个统计还是有点意思，说下它的两种解法。

Treeloveswater的点分治

.

往点分治方向思考，问题就变成了：你有一棵树，如何\(O(n)\)的处理出，以根为lca的点对的答案？

一个很重要的性质：

对于树中的一点i，如果该点的颜色在该点到根这条链上是第一次出现，那么对于这棵树的其他与i的lca为根点j（即在不同子树内），均能与i的子树（包括i）组成点对，i的颜色会对j的答案贡献size[i]。（我们在此暂且不考虑j到根的链上是否出现了i的颜色，待会儿容斥掉）

这个性质很显然。

那么我们就可以这样做了：

1. 对树进行第一遍dfs，预处理size和上方性质中每个颜色的贡献color，同时记录color总和sum

2. 枚举根的所有儿子子树，先把子树扫一遍清除其在color数组中的所有贡献（排除同一子树内部的错误贡献）。接着，对于该子树中的每一个点j：
   
   设X=sigma color[j 到根上（不包括根）的所有颜色] （由于这些颜色已经出现过，我们不能在该子树外计算其贡献）
   
   设num为j到根上（不包括根）的颜色数
   
   设Ｙ为size[root]-size[该子树（注意不是j）]（即所有其他子树＋根的点数）
   
   则ans[j]+=sum-X+num*Y

3. 别忘了计算root的ans
   
   ans[root]+=sum-color[根的颜色]+size[root]

那么点分治就解决了这个问题，时间复杂度\(O(n\log n)\)。统计方法值得学习。

看一下别人的代码。

#include<cstdio>
#include<cstring>
#include<iostream>
#include<algorithm>
#define o 200011
#define ll long long
using namespace std;
const int inf=1e8;
int head[o],nxt[o*2],point[o*2],V[o];
ll color[o],ans[o],much,sum,num,size[o],cnt[o],total,record;
int tot,n,ui,vi,root;
bool vis[o*2];
void addedge(int x,int y){
    tot++;nxt[tot]=head[x];head[x]=tot;point[tot]=y;
    tot++;nxt[tot]=head[y];head[y]=tot;point[tot]=x;
}
void findroot(int now,int dad){
    size[now]=1;
    ll maxson=0;
    for(int tmp=head[now];tmp;tmp=nxt[tmp]){
        int v=point[tmp];
        if(v!=dad&&!vis[tmp]){
            findroot(v,now);
            size[now]+=size[v];
            maxson=max(maxson,size[v]);
        }
    }
    maxson=max(maxson,total-size[now]);
    if(maxson<record) root=now,record=maxson;
}
void dfs1(int now,int dad){
    size[now]=1;
    cnt[V[now]]++;
    for(int tmp=head[now];tmp;tmp=nxt[tmp]){
        int v=point[tmp];
        if(!vis[tmp]&&v!=dad){
            dfs1(v,now);
            size[now]+=size[v];
        }
    }
    if(cnt[V[now]]==1){
        sum+=size[now];
        color[V[now]]+=size[now];
    }
    cnt[V[now]]--;
}
void change(int now,int dad,int value){
    cnt[V[now]]++;
    for(int tmp=head[now];tmp;tmp=nxt[tmp]){
        int v=point[tmp];
        if(!vis[tmp]&&v!=dad) change(v,now,value);
    }
    if(cnt[V[now]]==1){
        sum+=(ll)size[now]*value;
        color[V[now]]+=(ll)size[now]*value;
    }
    cnt[V[now]]--;
}
void dfs2(int now,int dad){
    cnt[V[now]]++;
    if(cnt[V[now]]==1){
        sum-=color[V[now]];
        num++;
    }
    ans[now]+=sum+num*much;
    for(int tmp=head[now];tmp;tmp=nxt[tmp]){
        int v=point[tmp];
        if(!vis[tmp]&&v!=dad) dfs2(v,now);
    }
    if(cnt[V[now]]==1){
        sum+=color[V[now]];
        num--;
    }
    cnt[V[now]]--;
}
void clear(int now,int dad){
    cnt[V[now]]=0;
    color[V[now]]=0;
    for(int tmp=head[now];tmp;tmp=nxt[tmp]){
        int v=point[tmp];
        if(!vis[tmp]&&v!=dad) clear(v,now);
    }
}
void solve(int now,int dad){
    dfs1(now,dad);
    ans[now]+=sum-color[V[now]]+size[now];
    for(int tmp=head[now];tmp;tmp=nxt[tmp]){
        int v=point[tmp];
        if(!vis[tmp]&&v!=dad){
            cnt[V[now]]++;
            sum-=size[v];
            color[V[now]]-=size[v];
            change(v,now,-1);
            cnt[V[now]]--;
            much=size[now]-size[v];
            dfs2(v,now);
            cnt[V[now]]++;
            sum+=size[v];
            color[V[now]]+=size[v];
            change(v,now,1);
            cnt[V[now]]--;
        }
    }
    sum=0;num=0;
    clear(now,dad);
    for(int tmp=head[now];tmp;tmp=nxt[tmp]){
        int v=point[tmp];
        if(!vis[tmp]&&v!=dad){
            vis[tmp]=true;
            vis[tmp^1]=true;
            total=size[v];
            record=inf;
            findroot(v,now);
            solve(root,0);
        }
    }
}
int main(){
    tot=1;
    scanf("%d",&n);
    for(int i=1;i<=n;i++) scanf("%d",&V[i]);
    for(int i=1;i<n;i++){
        scanf("%d%d",&ui,&vi);
        addedge(ui,vi);
    }
    record=inf;
    total=n;
    findroot(1,0);
    solve(root,0);
    for(int i=1;i<=n;i++) printf("%lld\n",ans[i]);
    return 0;
}

sxd666888的树上差分

分开计算每种颜色对答案的贡献。

我们考虑把树中这种颜色的点都删掉，那么就会有很多的小树，这些小树中的点互相之间不会产生贡献，而不同树的两个点之间会产生贡献。我们可以得到点的sum要+=n - 所在小树的size。

因此，一个点的sum=n * 颜色数 - 计算每种颜色节点时该点所在小树的size。发现我们只需要计算减号后的部分。

考虑在每棵小树的树根（深度最小）计算这棵小树的size，这样既方便计算也方便向下传递。我们用surp[i]记录把fa对应颜色删掉后i所在小树（i一定是这棵小树的树根）的size。

如何算所有颜色对一个点的贡献总和呢？直接维护总和sum，考虑在i时继承总和sum，把sum加上surp[i]，减去上一次同一颜色的surp更新，就行了。

特殊处理一下整棵树的根节点就好了。时间复杂度\(O(n)\)。

看一下此人的毒瘤命名代码。

#include<bits/stdc++.h>
using namespace std;
long long read()
{
    char ch=getchar();long long x=0,ff=1;
    while(ch<'0'||ch>'9') {if(ch=='-') ff=-1;ch=getchar();}
    while(ch>='0'&&ch<='9') x=x*10+ch-'0',ch=getchar();
    return x*ff;
}
void write(long long aa)
{
    if(aa<0) putchar('-'),aa=-aa;
    if(aa>9) write(aa/10);
    putchar('0'+aa%10);
    return;
}
long long n,sum,qwq;
long long vis[100005],ans[100005];
long long tot,head[100005],nx[200005],to[200005];
long long col[100005],sz[100005],jian[100005]；
long long lz[100005],bj[100005];
void jia(long long aa,long long bb)
{
    tot++;
    nx[tot]=head[aa];
    to[tot]=bb;
    head[aa]=tot;
    return;
}
void dfs(long long rt,long long fa)
{
    sz[rt]=1;
    long long tmp=jian[col[fa]];//遍历时删的个数
    for(long long i=head[rt];i;i=nx[i])
    {
        long long yy=to[i];
        if(yy==fa) continue;
        dfs(yy,rt);
        sz[rt]+=sz[yy];
    }
    jian[col[rt]]++;//删点
    if(fa)
    {
        lz[rt]=sz[rt]-jian[col[fa]]+tmp;//子树的size - （当前删的个数 - 遍历时删的个数）
        jian[col[fa]]+=lz[rt];//删点
    }
}
void getans(long long rt,long long fa)
{
    long long yuanbj=bj[col[fa]];
    qwq+=lz[rt]-bj[col[fa]];//差分啦
    bj[col[fa]]=lz[rt];
    ans[rt]=n*sum-qwq+bj[col[rt]];//自己颜色的显然是不能删掉的
    for(long long i=head[rt];i;i=nx[i])
    {
        long long yy=to[i];
        if(yy==fa) continue;
        getans(yy,rt);
    }
    bj[col[fa]]=yuanbj;
    qwq-=lz[rt]-bj[col[fa]];//还原啦
    return;
}
int main()
{
    n=read();
    for(long long i=1;i<=n;++i)
    {
        col[i]=read();//col[i]<=100000,可能大于n。。。。
        if(!vis[col[i]]) vis[col[i]]=1,sum++;//sum颜色种类
    }
    for(long long i=1;i<n;++i)
    {
        long long x=read(),y=read();
        jia(x,y);jia(y,x);
    }
    dfs(1,0);
    for(long long i=1;i<=100000;++i) if(vis[i]) qwq+=n-jian[i],bj[i]=n-jian[i];//特别处理根节点
    getans(1,0);
    for(long long i=1;i<=n;++i) write(ans[i]),puts("");
    return 0;
}