P2664 树上游戏

题目

洛谷

做法

刚学完点分治，"点分治真好玩"，\(dalao\)瞬间把这题丢过来，肛了三个小时终于\(A\)了

\[ans_i=\sum\limits_{j=1}^n sum(i,j) \]

模板题告诉我们：要把其他子树都堆到一起进行计算，这题相对来说并不好处理

但其实是具有单调性的，设树的重心为\(w\)，一对父子关系的点对\((u,v)\)，\(v\)是包含\(u\)的

更新：将重心看作一个无色节点，遍历其中一棵子树，能得到子树总贡献及每种颜色的贡献

统计答案：点对\((u,v)\)，\(v\)统计在其他子树的贡献一定比\(u\)小，之前已得出其他子树每种颜色的贡献，减掉就好了

细节：

1. 统计重心贡献
2. 统计每种颜色贡献的标记记得弹出子树后清空
3. 实际菊花图能卡爆空间，可以用\(map\)处理数组，由于本题数据过水开子树序列\(100\)就能过了

My complete code

#include<bits/stdc++.h>
using namespace std;
typedef long long LL;
inline LL Read(){
    LL x(0),f(1); char c=getchar();
    while(c<'0' || c>'9'){ if(c=='-')f=-1; c=getchar(); }
    while(c>='0' && c<='9') x=(x<<3)+(x<<1)+c-'0',c=getchar();
    return x*f;
}
const LL maxn=1e5+9,inf=0x3f3f3f3f;
struct node{
    LL to,nxt;
}dis[maxn<<1];
void Clear(LL u,LL now);
LL num,root,N,mi,tmp,n,tree,nsize;
LL head[maxn],size[maxn],a[maxn],cnt[100][maxn],ans[maxn],fir[maxn],V[maxn],sz[100],nsum[100],st[maxn];
bool visit[maxn];
inline void Add(LL u,LL v){
    dis[++num]=(node){v,head[u]}, head[u]=num;
}
void Dfs(LL u){
    size[u]=1; LL mx(0); visit[u]=true;
    for(LL i=head[u];i;i=dis[i].nxt){
        LL v(dis[i].to);
        if(visit[v]) continue;
        Dfs(v), size[u]+=size[v];
        mx=max(mx,size[v]);
    }mx=max(mx,N-size[u]);
    if(mi>mx) mi=mx,root=u;
    visit[u]=false;
}
void Up(LL u,LL now){
    ++sz[0], ++sz[now];
    visit[u]=true;
    bool f(false);
    if(!fir[a[u]]) 
        fir[a[u]]=f=true,
        cnt[0][a[u]]+=size[u],cnt[now][a[u]]+=size[u],
        nsum[0]+=size[u],nsum[now]+=size[u];
    for(LL i=head[u];i;i=dis[i].nxt){
        LL v(dis[i].to);
        if(visit[v]) continue;
        Up(v,now);
    }visit[u]=false;
    if(f) fir[a[u]]=false;
}
void Get(LL u,LL now){
    bool f(false);
    if(!fir[a[u]])
        fir[a[u]]=f=true,tmp=tmp-cnt[0][a[u]]+cnt[now][a[u]],++tree;
    ans[root]+=tree;
    ans[u]+=tmp+tree*nsize; visit[u]=true;
    for(LL i=head[u];i;i=dis[i].nxt){
        LL v(dis[i].to);
        if(visit[v]) continue;
        Get(v,now);
    }visit[u]=false;
    if(f) --tree,fir[a[u]]=false,tmp=tmp+cnt[0][a[u]]-cnt[now][a[u]];
}
void Div(LL u){
    mi=inf, Dfs(u);
    Dfs(root);
    visit[u=root]=true;
    memset(nsum,0,sizeof(nsum));
    nsize=1;
    LL now(0);
    for(LL i=head[u];i;i=dis[i].nxt){
        LL v(dis[i].to);
        if(visit[v]) continue;
        ++now;
        Up(v,now);
    }now=0;
    for(LL i=head[u],S=nsize;i;i=dis[i].nxt){
        LL v(dis[i].to); 
        if(visit[v]) continue;
        ++now;
        nsize=sz[0]-sz[now]+1;
        tmp=nsum[0]-nsum[now];
        fir[a[u]]=true, tree=1;
        tmp=tmp-cnt[0][a[u]]+cnt[now][a[u]];
        Get(v,now);
        fir[a[u]]=false;
    }
    now=0;
    for(LL i=head[u];i;i=dis[i].nxt){
        LL v(dis[i].to); if(visit[v]) continue;
        ++now;
        Clear(v,now);
    }
    
    for(LL i=head[u],sum=N;i;i=dis[i].nxt){
        LL v(dis[i].to);
        if(visit[v]) continue;
        N=size[v];
        Div(v);
    }
}
void Clear(LL u,LL now){
    visit[u]=true;
    cnt[0][a[u]]=cnt[now][a[u]]=0;
    sz[0]=sz[now]=0;
    nsum[0]=nsum[now]=0;
    for(LL i=head[u];i;i=dis[i].nxt){
        LL v(dis[i].to);
        if(visit[v]) continue;
        Clear(v,now);
    }visit[u]=false;
}
int main(){
    n=Read();
    for(LL i=1;i<=n;++i) a[i]=Read();
    for(LL i=1;i<n;++i){
        LL u(Read()),v(Read());
        Add(u,v), Add(v,u);
    }
    N=n, Div(1);
    for(LL i=1;i<=n;++i) printf("%lld\n",ans[i]+1);
    return 0;
}