Codechef:Path Triples On Tree

Path Triples On Tree

题意是求树上都不相交或者都相交的路径三元组数量。

发现blog里没什么树形dp题，也没有cc题，所以来丢一道cc上的树形dp题。

比较暴力，比较恶心

#include<cstdio>
#include<algorithm>
#define MN 300001
#define ui long long
using namespace std;

ui read_p,read_ca;
inline ui read(){
    read_p=0;read_ca=getchar();
    while(read_ca<'0'||read_ca>'9') read_ca=getchar();
    while(read_ca>='0'&&read_ca<='9') read_p=read_p*10+read_ca-48,read_ca=getchar();
    return read_p;
}
ui n,N,m,l[MN],num=0,fa[MN],s[MN],d[MN],_s[MN],S[MN],GS[MN],_S[MN],si[MN],MMH=0,_MMH=0,x,y;
struct na{ui y,ne;}b[MN<<1];
inline void in(ui x,ui y){b[++num].y=y;b[num].ne=l[x];l[x]=num;}
const ui MOD=1e9+7;
inline ui _M(ui x){while(x>=MOD)x-=MOD;while(x<0)x+=MOD;return x;}
inline void M(ui &x){while(x>=MOD)x-=MOD;while(x<0)x+=MOD;}
void dfs(ui x){
    si[x]=1;
    ui i,k=0,al=0;
    for (i=l[x];i;i=b[i].ne) if (!si[b[i].y]) fa[b[i].y]=x,dfs(b[i].y),al=(1LL*si[x]*si[b[i].y]+al)%MOD,si[x]+=si[b[i].y],M(s[x]+=s[b[i].y]),M(k+=d[b[i].y]);
    
    for (i=l[x];i;i=b[i].ne) if (b[i].y!=fa[x])
    M(MMH+=1LL*d[b[i].y]*(al-1LL*si[b[i].y]*(si[x]-si[b[i].y])%MOD)%MOD+1LL*_s[b[i].y]*(si[x]-si[b[i].y])%MOD-MOD),M(GS[x]+=1LL*s[b[i].y]*(si[x]-si[b[i].y])%MOD+GS[b[i].y]-MOD);
    
    s[x]=(1LL*al*si[x]+s[x])%MOD;
    M(S[x]=(1LL*(si[x]-1)*si[x]>>1)%MOD*si[x]%MOD-s[x]);
    d[x]=(1LL*al*(al-1)>>1)%MOD;
    _s[x]=k;
    
    for (i=l[x];i;i=b[i].ne) if (b[i].y!=fa[x])
    d[x]=(1LL*(si[x]-si[b[i].y])*s[b[i].y]+d[b[i].y]+d[x])%MOD,_s[x]=(1LL*(k-d[b[i].y])*si[b[i].y]+_s[x]+_s[b[i].y])%MOD
    ,M(_S[x]+=(1LL*GS[b[i].y]*(si[x]-si[b[i].y])+(-1LL*si[b[i].y]*(si[x]-si[b[i].y])%MOD+al)*s[b[i].y]+_S[b[i].y])%MOD);
    
}
void DFS(ui x,ui v,ui _v){
    ui i,S=0,AL=0,k=0,_k=0,P=0,al=0,o;
    for (i=l[x];i;i=b[i].ne) if (b[i].y!=fa[x])
    M(al+=1LL*(si[x]-si[b[i].y])*si[b[i].y]%MOD),M(k+=s[b[i].y]),M(_k+=1LL*s[b[i].y]*(n-si[b[i].y])%MOD),
    M(P+=d[b[i].y]),M(AL+=(1LL*(si[b[i].y]+1)*(n-si[b[i].y])-1)%MOD*si[b[i].y]%MOD);
    al=1LL*(al+si[x]-1)*((MOD+1)/2)%MOD;
    M(MMH+=1LL*v*al%MOD);
    
    for (i=l[x];i;i=b[i].ne) if (b[i].y!=fa[x]) o=S=_M((1LL*(si[x]-si[b[i].y])*(n-si[x])-1LL*(si[x]-si[b[i].y])*si[b[i].y]+al)%MOD),o=1LL*o*(n-si[b[i].y])%MOD,S=(1LL*S*(S-1)>>1)%MOD,
    M(S+=P-d[b[i].y]),DFS(b[i].y,_M(_M(v+_k-1LL*s[b[i].y]*(n-si[b[i].y])%MOD-1LL*(k-s[b[i].y])*si[b[i].y]%MOD)+1LL*(si[x]-si[b[i].y])*_v%MOD-MOD+S),_M(_v+k-MOD-s[b[i].y]+o));
}
void _DFS(ui x){
    ui i,k=0,al=0;
    for (i=l[x];i;i=b[i].ne) if (b[i].y!=fa[x]) _DFS(b[i].y),M(al+=1LL*(si[x]-si[b[i].y])*si[b[i].y]%MOD);
    al=1LL*(al+si[x]-1)*((MOD+1)/2)%MOD;
    for (i=l[x];i;i=b[i].ne) if (b[i].y!=fa[x]) M(_MMH+=((1LL*(si[x]-si[b[i].y])*(n-si[x])-1LL*(si[x]-si[b[i].y])*si[b[i].y]+al)%MOD*(si[x]-si[b[i].y])%MOD+k)*s[b[i].y]%MOD),M(k+=s[b[i].y]);
    for (i=l[x];i;i=b[i].ne) if (b[i].y!=fa[x]) M(_MMH+=1LL*GS[b[i].y]*(n-si[b[i].y])%MOD*(si[x]-si[b[i].y])%MOD),M(_MMH+=1LL*_S[b[i].y]*(si[x]-si[b[i].y])%MOD);
}
int main(){
    register ui i;
    n=read();
    for (i=1;i<n;i++) x=read(),y=read(),in(x,y),in(y,x);
    N=(1LL*n*(n-1)>>1)%MOD;m=1LL*N*(N-1)%MOD*(N-2)%MOD*((MOD+1)/6)%MOD;
    dfs(1);
    DFS(1,0,0);
    _DFS(1);
    M(m-=MMH+_MMH);
    printf("%lld
",m);
}