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\n",m);
}

?