[luogu2664]树上游戏（点分治）

题目：

我是超链接

题解：

这种树上路径相关的考虑点分治，对于当前的分治中心，需要统计出【自己出发到分治块内所有路径对自己答案的贡献】和【经过ta的路径对当前分治块内点的贡献】。这个自己出发到所有点的好求，问题是第二个怎么求

我们对于当前的分治中心的每一个子树考虑，令cnt[i]表示从分治中心出发的所有路径包括颜色i的条数，size为分治块所有点的个数。那么我们在dfs到x的时候，首先给sum[x]增加$\sum cnt[i]$，即所有在其他子树中出现过的颜色总和，然后计算x到分治中心路径上的颜色贡献

对于一个出现在分治中心到x的路径上的颜色c，他对x的贡献为size-cnt[c]，因为c已经在一些路径上出现，他能产生额外的贡献就是ta原来没有出现过的路径条数，所以还要给sum[x]加上size-cnt[c]。同时c会对x的子树内所有点产生贡献，这个贡献要像标记一样往下传递，然后标记一下c的贡献已经计算过不用再算了

所以只需要对当前分治中心求出cnt和size，进入一棵子树时减去子树自己内部对cnt产生的贡献。为了保证时间复杂度，我们不能对于所有颜色求cnt，要先统计当前分治块内哪些颜色出现了，这样枚举块内所有颜色的复杂度才是O(分治块)大小

总结一下，这道题是点分治，整体的思路是在O(n)的时间内处理出以分治中心为lca的子树内的两个点的贡献
第一部分是排除自己子树后的Σcnt，第二部分是dfs的时候如果这个颜色i第一次出现，要加上size[root]-cnt[i]（并给子树加上标记）

代码：

#include <cstdio>
#include <iostream>
#define LL long long
#define INF 1e9
using namespace std;
const int N=100005;
int tot,nxt[N*2],v[N*2],point[N],a[N],size[N],f[N],have[N],col[N],copycol[N],root,jd,num,n;bool vis[N],ext[N];
LL sum[N],cnt[N],all,S,copycnt[N],soncnt[N];
void addline(int x,int y)
{
    ++tot; nxt[tot]=point[x]; point[x]=tot; v[tot]=y;
    ++tot; nxt[tot]=point[y]; point[y]=tot; v[tot]=x;
}
void find(int x,int fa)
{
    size[x]=1; f[x]=0;
    for (int i=point[x];i;i=nxt[i])
      if (v[i]!=fa && !vis[v[i]])
      {
        find(v[i],x);
        size[x]+=size[v[i]];
        f[x]=max(f[x],size[v[i]]);
      }
    f[x]=max(f[x],jd-f[x]);
    if (f[x]<f[root]) root=x;
}
void findcnt(int x,int fa,LL *cnt)
{
    if (!ext[a[x]]) ext[a[x]]=1,col[++num]=a[x];
    if (++have[a[x]]==1) cnt[a[x]]+=size[x];
    for (int i=point[x];i;i=nxt[i])
      if (v[i]!=fa && !vis[v[i]]) findcnt(v[i],x,cnt);
    have[a[x]]--;
}
void change(int x,int fa,LL last)
{
    if (++have[a[x]]==1) last+=all-copycnt[a[x]];
    sum[x]+=last+S;
    for (int i=point[x];i;i=nxt[i])
      if (v[i]!=fa && !vis[v[i]]) change(v[i],x,last);
    have[a[x]]--;
}
void calc(int x)
{
    find(x,0);
    //必不可少的一步，更改size关系 
    S=num=0;int Num;
    findcnt(x,0,cnt);
    Num=num;
    for (int i=1;i<=num;i++) ext[copycol[i]=col[i]]=0,S+=cnt[col[i]],copycnt[col[i]]=cnt[col[i]];
    sum[x]+=S;LL SS=S;
    for (int i=point[x];i;i=nxt[i])
      if (!vis[v[i]])
      {
        have[a[x]]=1;
        num=0;
        findcnt(v[i],x,soncnt);
        have[a[x]]=0;
        S-=size[v[i]]; copycnt[a[x]]-=size[v[i]];
        //不管怎么样，a[x]这个点的颜色在ta的子树中肯定有一个贡献 
        //所以我们既然要排除这棵子树的影响，理应如此 
        for (int j=1;j<=num;j++) ext[col[j]]=0,copycnt[col[j]]-=soncnt[col[j]],S-=soncnt[col[j]];

        all=size[x]-size[v[i]];
        change(v[i],x,0);

        S=SS; copycnt[a[x]]+=size[v[i]];
        for (int j=1;j<=num;j++) copycnt[col[j]]=cnt[col[j]],soncnt[col[j]]=0;
      }
    for (int i=1;i<=Num;i++) cnt[copycol[i]]=0;
}
void work(int x)
{
    calc(x);vis[x]=1;
    for (int i=point[x];i;i=nxt[i])
      if (!vis[v[i]])
      {
        jd=size[v[i]]; f[0]=INF; root=0; find(v[i],x);
        work(root);
      }
}
int main()
{
    scanf("%d",&n);
    for (int i=1;i<=n;i++) scanf("%d",&a[i]);
    for (int i=1;i<n;i++)
    {
        int x,y;scanf("%d%d",&x,&y);
        addline(x,y);
    }
    jd=n; f[0]=INF; root=0; find(1,0);
    work(root);
    for (int i=1;i<=n;i++) printf("%lld\n",sum[i]);
}