[点分治] Luogu P2664 树上游戏

题解

• 对于树中的一点i，如果该点的颜色在该点到根这条链上是第一次出现，那么对于这棵树的其他点j（以i和j的lca为根），均能与ｉ的子树（包括i）组成点对，i的颜色会对j的答案贡献size[i]
• 然后点分治就可以解决了

代码

#include <cstdio>
#include <iostream>
#define ll long long
using namespace std;
const int N=2e5+10,inf=1e8;
int n,cnt=1,root,head[N],v[N],vis[N*2];
ll q,sum,num,tot,Mx=inf,size[N],col[N],ans[N],p[N];
struct edge { int to,from; }e[N*2];
void clear(int x,int fa) { p[v[x]]=col[v[x]]=0; for (int i=head[x];i;i=e[i].from) if (!vis[i]&&e[i].to!=fa) clear(e[i].to,x); }
void insert(int x,int y)
{
    e[++cnt].to=y,e[cnt].from=head[x],head[x]=cnt;
    e[++cnt].to=x,e[cnt].from=head[y],head[y]=cnt;
}
void dfs1(int x,int fa)
{
    size[x]=1,p[v[x]]++;
    for (int i=head[x];i;i=e[i].from) if (e[i].to!=fa&&!vis[i]) dfs1(e[i].to,x),size[x]+=size[e[i].to];
    if (p[v[x]]==1) sum+=size[x],col[v[x]]+=size[x];
    p[x[v]]--;
}
void dfs2(int x,int fa)
{
    p[v[x]]++;
    if (p[v[x]]==1) sum-=col[v[x]],num++;
    ans[x]+=sum+num*q;
    for (int i=head[x];i;i=e[i].from) if (!vis[i]&&e[i].to!=fa) dfs2(e[i].to,x);
    if (p[v[x]]==1) sum+=col[v[x]],num--;
    p[v[x]]--;
}
void findroot(int x,int fa)
{
    size[x]=1; ll mx=0;
    for (int i=head[x];i;i=e[i].from) if (e[i].to!=fa&&!vis[i]) findroot(e[i].to,x),size[x]+=size[e[i].to],mx=max(mx,size[e[i].to]);
    mx=max(mx,tot-size[x]);
    if (mx<Mx) Mx=mx,root=x;
}
void change(int x,int fa,int k)
{
    p[v[x]]++;
    for (int i=head[x];i;i=e[i].from) if (e[i].to!=fa&&!vis[i]) change(e[i].to,x,k);
    if (p[v[x]]==1) sum+=(ll)size[x]*k,col[v[x]]+=(ll)size[x]*k;
    p[v[x]]--;
}
void solve(int x,int fa)
{
    dfs1(x,fa),ans[x]+=sum-col[v[x]]+size[x];
    for (int i=head[x];i;i=e[i].from) if (!vis[i]&&e[i].to!=fa)
        p[v[x]]++,sum-=size[e[i].to],col[v[x]]-=size[e[i].to],change(e[i].to,x,-1),p[v[x]]--,q=size[x]-size[e[i].to],dfs2(e[i].to,x),p[v[x]]++,sum+=size[e[i].to],col[v[x]]+=size[e[i].to],change(e[i].to,x,1),p[v[x]]--;
    sum=0,num=0,clear(x,fa);
    for (int i=head[x];i;i=e[i].from) if (!vis[i]&&e[i].to!=fa) vis[i]=1,vis[i^1]=1,tot=size[e[i].to],Mx=inf,findroot(e[i].to,x),solve(root,0);
}
int main()
{
    scanf("%d",&n);
    for (int i=1;i<=n;i++) scanf("%d",&v[i]);
    for (int i=1,x,y;i<n;i++) scanf("%d%d",&x,&y),insert(x,y);
    tot=n,findroot(1,0),solve(root,0);
    for (int i=1;i<=n;i++) printf("%lld
",ans[i]);
}