树上游戏（树上乱搞）

树上游戏（luogu）

Solution

• 对于每种颜色 i ，计算出以每个点为起点的包含它的路径条数，每个点的答案为各种颜色加起来
• 删去颜色为 i 的所有点，树变成森林
• 每颗树内点间的路径不包含颜色 i ，不能对包含颜色 i 的路径条数产生贡献
• 不同树间的点，路径上必然有颜色 i 
• 即颜色不为 i 的点，以它为起点的包含颜色 i 的路径条数为——n-删去颜色为 i 的所有点后它所在树的节点数
• 颜色为 i 的点，以它为起点的包含颜色 i 的路径条数为 n
• 第一次遍历，求出每个点去掉父亲颜色后所在树的节点数（该数以这个点为根）
• 特殊处理原树的根（它没有父亲，可看做去掉任何颜色后树的根）
• 第二次遍历，差分求出答案

Code

#include <cstdio>
#include <cstdlib>
#define ll long long
using namespace std;
const int N=1e5+10;
int head[N],nxt[N*2],ver[N*2],tot;
int si[N],tr[N],cut[N],n,col[N],u,v,d[N],cnt[N];
bool flag[N];
ll ans[N],sum;
void add(int x,int y)
{
    ver[++tot]=y,nxt[tot]=head[x],head[x]=tot;
}
void dfs(int u,int fa)
{
    int pre=cut[col[fa]];
    si[u]=1;
    for(int i=head[u];i;i=nxt[i])
    {
        int v=ver[i];
        if(fa==v) continue;
        dfs(v,u);
        si[u]+=si[v];
    }
    cut[col[u]]++;
    if(!fa) return ;
    tr[u]=si[u]-(cut[col[fa]]-pre);
    cut[col[fa]]+=tr[u];
}
void solve(int u,int fa)
{
    int pre=d[col[fa]];
    d[col[fa]]=tr[u];
    sum+=tr[u]-pre;
    ll aa=n,bb=cnt[0],cc=sum,dd=d[col[u]];
    ans[u]=aa*bb-cc+dd;
    for(int i=head[u];i;i=nxt[i])
    {
        int v=ver[i];
        if(v==fa) continue;
        solve(v,u);
    }
    sum-=tr[u]-pre;
    d[col[fa]]=pre;
}
int main()
{
    scanf("%d",&n);
    for(int i=1;i<=n;i++)
    {
        scanf("%d",&col[i]);
        if(!flag[col[i]]) flag[col[i]]=true,cnt[++cnt[0]]=col[i];
    }
    for(int i=1;i<n;i++)
    {
        scanf("%d%d",&u,&v);
        add(u,v),add(v,u);
    }
    dfs(1,0);
    for(int i=1;i<=cnt[0];i++) d[cnt[i]]=n-cut[cnt[i]],sum+=d[cnt[i]];
    solve(1,0);
    for(int i=1;i<=n;i++) printf("%lld\n",ans[i]);
    return 0;
}

 

以它为起点的包含颜色 i 的路径条数为