洛谷 2664 树上游戏——点分治 / 思路

题目：https://www.luogu.org/problemnew/show/P2664

想到点分治的话，想到一条路径会对路径的两端点的 ans[ ] 有贡献，所以每次就考虑过重心的路径。

先想的是枚举颜色，每次做一种颜色。

枚举重心的一个子树的时候，就记录一下当前重心 cr 的其它子树里有几条 “从 cr 开始，含有该颜色” 的路径，设为 sc；

在当前路径上如果没碰见有颜色的点，那么给当前点 ans[ ] += sc ；如果碰见过有颜色的点，就给当前点 ans[ ] += siz(其他子树)+1 （+1表示重心那个点）；

先把所有子树的 sc 累加起来，进入这棵子树之前原样 dfs 一下该子树把它带来的 sc 减去，再进入它算答案即可。

然后考虑所有颜色也可以一起做。就是把 sc 扩充成一个数组，用 vc[ i ] 表示当前路径有没有曾经有过 i 这种颜色（若使 vc = 1 的点是 cr ，则 use[ cr ] = 1 ，方便回溯的时候把 vc 改回来）。

先把所有颜色的 sc（其他子树） 累加起来，记为 s1 ；如果当前路径碰到一种颜色 i ，就在 s1 里减掉 sc[ i ]（其他子树），而给一个初始为 0 的 lj += siz（其他子树），表示该颜色的贡献不要求 “其他子树路径有该颜色” 了。

注意重心那个点也对 vc 和 sc 有影响。并注意最后删除影响的时候记得删除重心带来的影响！

#include<cstdio>
#include<cstring>
#include<algorithm>
#define ll long long
using namespace std;
int rdn()
{
  int ret=0;bool fx=1;char ch=getchar();
  while(ch>'9'||ch<'0'){if(ch=='-')fx=0;ch=getchar();}
  while(ch>='0'&&ch<='9')ret=ret*10+ch-'0',ch=getchar();
  return fx?ret:-ret;
}
int Mn(int a,int b){return a<b?a:b;}
int Mx(int a,int b){return a>b?a:b;}
const int N=1e5+5,mod=1e9+7;
int n,c[N],siz[N],hd[N],xnt,to[N<<1],nxt[N<<1],rt,mn;
ll ans[N],sc[N],s1,s2; bool vis[N],use[N],vc[N];
void add(int x,int y){to[++xnt]=y;nxt[xnt]=hd[x];hd[x]=xnt;}
void get_rt(int cr,int fa,int s)
{
  siz[cr]=1; int mx=0;
  for(int i=hd[cr],v;i;i=nxt[i])
    if(!vis[v=to[i]]&&v!=fa)
      {
    get_rt(v,cr,s); siz[cr]+=siz[v];
    mx=Mx(mx,siz[v]);
      }
  mx=Mx(mx,s-siz[cr]); if(mx<mn)mn=mx,rt=cr;
}
void ini_dfs(int cr,int fa)
{
  siz[cr]=1;
  for(int i=hd[cr],v;i;i=nxt[i])
    if(!vis[v=to[i]]&&v!=fa)
      ini_dfs(v,cr), siz[cr]+=siz[v];
}
void dfs(int cr,int fa,int fx)
{
  int nw=c[cr];
  if(!vc[nw])
    {
      vc[nw]=1; use[cr]=1;
      int d=(fx?siz[cr]:-siz[cr]); sc[nw]+=d; s1+=d;
      if(fx==1)ans[rt]+=siz[cr];
    }
  for(int i=hd[cr],v;i;i=nxt[i])
    if(!vis[v=to[i]]&&v!=fa) dfs(v,cr,fx);
  if(use[cr]){vc[nw]=0; use[cr]=0;}
}
void dfsx(int cr,int fa,ll lj)
{
  int nw=c[cr];
  if(!vc[nw])
    { vc[nw]=1; use[cr]=1; s1-=sc[nw]; lj+=s2;}
  //printf(" dfsx:cr=%d s1=%lld lj=%lld\n",cr,s1,lj);
  ans[cr]+=s1+lj;
  for(int i=hd[cr],v;i;i=nxt[i])
    if(!vis[v=to[i]]&&v!=fa) dfsx(v,cr,lj);
  if(use[cr]){vc[nw]=0; use[cr]=0; s1+=sc[nw];}
}
void del_dfs(int cr,int fa)
{
  sc[c[cr]]=0;
  for(int i=hd[cr],v;i;i=nxt[i])
    if(!vis[v=to[i]]&&v!=fa) del_dfs(v,cr);
}
void solve(int cr,int s)
{
  //printf("solve:cr=%d (s=%d)\n",cr,s);
  vis[cr]=1; int nw=c[cr];
  vc[nw]=1; sc[nw]=s1=s; ans[rt]+=s;//s not s-1 for self
  for(int i=hd[cr],v;i;i=nxt[i])
    if(!vis[v=to[i]])ini_dfs(v,cr),dfs(v,cr,1);
  s2=s;
  for(int i=hd[cr],v;i;i=nxt[i])
    if(!vis[v=to[i]])
      {
    dfs(v,cr,0); sc[nw]-=siz[v]; s1-=siz[v]; s2=s-siz[v];
    /*printf(" v=%d(s1=%lld s2=%lld):",v,s1,s2);
      for(int j=1;j<=2;j++)printf("%lld ",sc[j]);puts("");*/
    dfsx(v,cr,0); sc[nw]+=siz[v]; s1+=siz[v];
    dfs(v,cr,2);//
      }
  /*printf("ans:");
    for(int i=1;i<=n;i++)printf("%lld ",ans[i]);puts("");*/
  vc[nw]=sc[nw]=0;///!!!
  for(int i=hd[cr],v;i;i=nxt[i])
    if(!vis[v=to[i]])del_dfs(v,cr);
  for(int i=hd[cr],v;i;i=nxt[i])
    if(!vis[v=to[i]])
      {
    mn=N;get_rt(v,0,siz[v]);solve(rt,siz[v]);
      }
}
int main()
{
  n=rdn();
  for(int i=1;i<=n;i++)c[i]=rdn();
  for(int i=1,u,v;i<n;i++)
    { u=rdn();v=rdn();add(u,v);add(v,u);}
  mn=N;get_rt(1,0,n);solve(rt,n);
  for(int i=1;i<=n;i++)printf("%lld\n",ans[i]);
  return 0;
}

看了题解，还有一种 O(n) 的做法。https://www.luogu.org/problemnew/solution/P2664

只看某种颜色的点，删掉它们，树变成一些 “小树” 。

该颜色对一个不是该颜色的点的贡献就是 n - siz[ 该点所在小树 ] 。

所以考虑在遍历树的时候把这样的 siz[ 小树 ] 都找出来。考虑一个某颜色下小树的 siz 记录在该小树最浅的节点上。这样 “该颜色” 一定是该点的父亲的颜色，且该点和父亲不同色。所以每个点只会记录一个信息。也就是想对每个点求出这部分的大小：（其中括住的那些点都是和父亲一个颜色）考虑令 sc[ i ] 表示该点子树里最浅点颜色为 i 的这种块的大小，则一个点记录的信息 f[ cr ] = siz[ cr ] - sc [ c[ fa ] ] （siz 是子树大小）。sc[ i ] 其实可以全局开一个，只要记录一下进 cr 子树之前的 sc 值，和遍历完 cr 子树的 sc 值对比一下就能得到子树里的 sc 值了。考虑统计答案的时候，如何知道这个点所在的其他颜色（非父亲颜色）下的“小树”大小？

其实这就是该点所有祖先的 f[ ] ，“父亲颜色” 相同的只保留最深的。这个只要开一个对于颜色的全局数组记录目前各颜色用的 f[ ] 值是多少就行了。

#include<cstdio>
#include<cstring>
#include<algorithm>
#define ll long long
using namespace std;
int rdn()
{
  int ret=0;bool fx=1;char ch=getchar();
  while(ch>'9'||ch<'0'){if(ch=='-')fx=0;ch=getchar();}
  while(ch>='0'&&ch<='9')ret=ret*10+ch-'0',ch=getchar();
  return fx?ret:-ret;
}
int Mx(int a,int b){return a>b?a:b;}
const int N=1e5+5;
int n,c[N],hd[N],xnt,to[N<<1],nxt[N<<1];
int siz[N],sc[N],f[N],ct; ll ans[N],sm; bool vis[N];
void add(int x,int y){to[++xnt]=y;nxt[xnt]=hd[x];hd[x]=xnt;}
void dfs(int cr,int fa)
{
  siz[cr]=1; int nw=c[fa], tp=sc[nw];
  for(int i=hd[cr],v;i;i=nxt[i])
    if((v=to[i])!=fa)
      dfs(v,cr), siz[cr]+=siz[v];
  sc[c[cr]]++; if(!fa)return;
  f[cr]=siz[cr]-(sc[nw]-tp);  sc[nw]=tp+siz[cr];
}
void dfsx(int cr,int fa)
{
  int nw=c[fa], tp=sc[nw];
  if(fa){ sm-=sc[nw]; sc[nw]=f[cr]; sm+=sc[nw];}
  ans[cr]=(ll)n*ct-sm+sc[c[cr]];
  for(int i=hd[cr],v;i;i=nxt[i])
    if((v=to[i])!=fa) dfsx(v,cr);
  if(fa){ sm-=sc[nw]; sc[nw]=tp; sm+=sc[nw];}
}
int main()
{
  n=rdn(); int mx=0;
  for(int i=1;i<=n;i++)
    {
      c[i]=rdn();mx=Mx(mx,c[i]);
      if(!vis[c[i]])vis[c[i]]=1, ct++;
    }
  for(int i=1,u,v;i<n;i++)
    {u=rdn();v=rdn();add(u,v);add(v,u);}
  dfs(1,0);
  for(int i=1;i<=mx;i++)if(vis[i])sc[i]=n-sc[i], sm+=sc[i];//vis!!!
  dfsx(1,0);
  for(int i=1;i<=n;i++)printf("%lld\n",ans[i]);
  return 0;
}