[51nod] 彩色树 树形dp

这道题是个人都会看得出来是一道树形Dp.

大概的思路就是计算每个点的贡献，也就是说计算有多少个点对经过该点，但事实上这样的计算非常麻烦，我们可以只计算不经过该点的点对数量（补集），这样子思路就非常显然了，据说可以优化到O(n），但吾辈太菜只能做O(nlogn)的

顺便吐槽一下，这道题是多校原题，而且名字还一样（出题人偷懒！）

 

#include<iostream>
#include<algorithm>
#include<cstdio>
#include<cstring>
#include<cmath>

using namespace std;
const int N = 100005,Mod = 1e9 + 7,Inv = (Mod+1)>>1;
typedef long long ll;
struct edge{
    int pos,nx;
}e[N*2];

ll res,n,num,size,head[N],color[N],son[N],cnt[N];
bool vis[N];

void insert(int u,int v){
    e[++num].pos=v;
    e[num].nx=head[u];
    head[u]=num;
}

void dfs(int x,int fa){
    int i;
    son[x]=1;
    ll z=cnt[color[x]],y=cnt[color[x]];
    for(i=head[x];i;i=e[i].nx){
        if(e[i].pos==fa) continue;
        dfs(e[i].pos,x);
        son[x]=(son[x]+son[e[i].pos])%Mod;
        ll temp=(son[e[i].pos]-cnt[color[x]]+z+Mod)%Mod;
        res=(res+((((temp-1ll+Mod)%Mod)*temp%Mod)*Inv)%Mod)%Mod;
        z=cnt[color[x]];
    }
    cnt[color[x]]=(son[x]+y)%Mod;
}

int main(){
    int x,y,i;
    scanf("%lld",&n);
    for(i=1;i<=n;i++) {
        scanf("%d",&color[i]);
        if(!vis[color[i]]) size++,vis[color[i]]=true;
    }
    
    for(i=1;i<n;i++){
        scanf("%d%d",&x,&y);
        insert(x,y);
        insert(y,x);
    }
    
    dfs(1,1);
    
    ll ans=(((1ll*n*(n-1))%Mod*size)%Mod*Inv)%Mod;
    for(i=1;i<=n;i++){
        if(vis[i]){
            ll temp=(n-cnt[i])%Mod;
            ans=(ans-(((temp*(temp-1))%Mod)*Inv)%Mod+Mod*2)%Mod;
        }
    }
    
    ans=(ans-res+Mod)%Mod;
    
    for(i=1;i<n;i++) ans=(ans*i)%Mod;
    
    printf("%lld",(ans*2)%Mod);
    return 0;
}