【JZOJ 3872】圣诞树

Description

圣诞节到了，小可可送给小薰一棵圣诞树。这棵圣诞树很奇怪，它是一棵多叉树，有n个点，n-1条边。它的每个结点都有一个权值。小可可和小薰想用这棵树玩一个游戏。
定义(s,e)为树上从s到e的简单路径，我们可以记下在这条路径上经过的结点，定义这个结点序列为S(s,e)。
我们按照如下方法定义这个序列S(s,e)的权值G(S(s,e))：假设这个序列中结点的权值为Z0，Z1，…，Z(L-1)，其中L为序列的长度，我们定义G(S(s,e))=Z0 × k^0 + Z1 × k^1 + … + Z(L-1) × k^(L-1)。
如果路径(s,e)满足G(S(s,e)) ≡ x (mod y) ，那么这条路径属于小可可，否则这条路径属于小薰。小可可和小薰很显然不希望这个游戏变得那么简单。小薰认为如果路径(p1,p2)和(p2,p3)都属于他，那么路径(p1,p3)也属于他，反之如果路径(p1,p2)和(p2,p3)都属于小可可，那么路径(p1,p3)也属于小可可。然而这个性质并不总是正确的。所以小薰想知道到底有多少三元组(p1,p2,p3)满足这个性质。
小薰表示她看一眼就知道这道题怎么做了。你会吗？

Solution

把所有点两两连边，如符合条件，则边权为1，反则为0，(可以自环)
那么，题目就变成了求多少个三角形，满足三条边边权均相同，
这个不太好做，试着减掉三条边不同的情况，
总共有一下6种情况：

一个点被指向1的次数为$in1$，指出1为$out1$，$in0,out0$类似，
发现，总的不合法的个数为$(in0*out1+in1*out0+2*(out1*out0+in0*in1))$
最后减的时候还要/2，因为每个被算了两遍，（注意：图中上2和上3两张图加起来为in0*in1，再加上下1,下3才乘2，这个读者可以自己画图算一下），

于是$O(n^2)$的算法出来了，
现在的问题是快速求那几个东西，
显然：$in0+in1=out0+out1=n$，所以只用计算其中的两个即可，
用点分治，每次计算出过重心的in1和out1，
观察计算的公式，写成先走到重心，再走到目的地的形式，合并一下，发现可以先移项，再把系数用逆元除过去，这样，整个试子就很清真，这里我用了哈希，
计算的时候可以用全局的减去本子树的，

细节还有很多，这里就不细讲了，

复杂度：$O(n\log(n))$

Code

#include <iostream>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#define fo(i,a,b) for(int i=a;i<=b;i++)
#define fod(i,a,b) for(int i=a;i>=b;i--)
#define efo(i,q) for(int i=A[q];i;i=B[i][0])
using namespace std;
typedef long long LL;
const int N=100500,MO=4965643;
int read(int &n)
{
    char ch=' ';int q=0,w=1;
    for(;(ch!='-')&&((ch<'0')||(ch>'9'));ch=getchar());
    if(ch=='-')w=-1,ch=getchar();
    for(;ch>='0' && ch<='9';ch=getchar())q=q*10+ch-48;n=q*w;return n;
}
int m,n;
LL X,mo,K,ans;
LL b[N];
int B[2*N][2],A[N],B0;
int IN[N],OUT[N];
bool z[N];
int cn[N],ct,ct1;
int H[4][MO][3],TI[2];//0:arrive 1:begin to go 2:ben arrive 3:ben begintogo
LL ny[N];
void link(int q,int w)
{
    B[++B0][0]=A[q];A[q]=B0,B[B0][1]=w;
    B[++B0][0]=A[w];A[w]=B0,B[B0][1]=q;
}
LL ksm(LL q,LL w)
{
    q=q%mo;LL ans=1;
    while(w)
    {
        if(w&1)ans=ans*q%mo;
        q=q*q%mo;w>>=1;
    }
    return ans;
}
int dfscn(int q,int fa,int ZN)
{
    cn[q]=1;
    int mx=0;
    efo(i,q)if(B[i][1]!=fa&&!z[B[i][1]])cn[q]+=dfscn(B[i][1],q,ZN),mx=max(mx,cn[B[i][1]]);
    mx=max(mx,ZN-cn[q]);
    if(mx<ct1)ct1=mx,ct=q;
    return cn[q];
}
int Ha(int I,LL q,bool jz)
{
    int i=q%MO,ti=TI[I>1];
    while(H[I][i][2]>=ti&&H[I][i][0]!=q)i=(i+1)%MO;
    if(H[I][i][2]<ti)H[I][i][1]=0,H[I][i][2]=ti,H[I][i][0]=q;
    if(jz)H[I][i][1]++;
    return H[I][i][1];
}
void dfsf(int q,int fa,LL cm,LL gt,LL k1,int jia,int c)
{
    cm=(cm+b[q]*k1%mo)%mo;Ha(0+jia,cm,1);
    gt=(gt*K%mo+b[q])%mo;Ha(1+jia,(X-gt+mo)%mo*ny[c+1]%mo,1);
    k1=k1*K%mo;
    efo(i,q)if(B[i][1]!=fa&&!z[B[i][1]])dfsf(B[i][1],q,cm,gt,k1,jia,c+1);
}
void dfs2(int q,int fa,LL cm,LL gt,LL k1,int c)
{
    cm=(cm+b[q]*k1%mo)%mo;IN[q]+=Ha(1,cm,0)-Ha(3,cm,0);
    gt=(gt*K+b[q])%mo;
    OUT[q]+=Ha(0,(X-gt+mo)%mo*ny[c+1]%mo,0)
        -Ha(2,(X-gt+mo)%mo*ny[c+1]%mo,0);
    k1=k1*K%mo;
    efo(i,q)if(B[i][1]!=fa&&!z[B[i][1]])dfs2(B[i][1],q,cm,gt,k1,c+1);
}
void divide(int q,int ZN)
{
    ct1=N;
    dfscn(q,0,ZN);
    q=ct;
    z[q]=1;TI[0]++;
    dfsf(q,0,0,0,1,0,0);
    IN[q]+=Ha(1,0,0);
    OUT[q]+=Ha(0,X,0);
    efo(i,q)if(!z[B[i][1]])
    {
        TI[1]++;
        dfsf(B[i][1],q,b[q],b[q],K,2,1);
        dfs2(B[i][1],q,0,0,1,0);
    }
    efo(i,q)if(!z[B[i][1]])divide(B[i][1],cn[B[i][1]]);
}
int main()
{
    int q,w;
    read(n);mo=read(q),K=read(q)%mo,X=read(w)%mo;
    LL t=1;
    fo(i,0,n)ny[i]=ksm(t,mo-2),t=t*K%mo;
    fo(i,1,n)b[i]=read(q)%mo;
    fo(i,2,n)read(q),read(w),link(q,w);
    divide(1,n);
    ans=0;
    fo(i,1,n)ans+=(n-IN[i])*OUT[i]+IN[i]*(n-OUT[i])+2*(OUT[i]*(n-OUT[i])+(n-IN[i])*IN[i]);
    printf("%lld\n",(LL)n*n*n-ans/2);
    return 0;
}