[联合集训6-9] El 点分治+哈希表

转化题意，就是根据题目所给的式子定义颜色，求同色三角形的个数。
有一个经典结论：完全图同色三角形个数=总三角形个数-异色角个数/2。因为每个异色三角形都恰有两个异色角。
在有向完全图中结论依然成立，但这里的异色角就要分别统计两出，一入一出，两出三种情况。我们只需要分别知道每个点进/出的红/绿路径数量即可。这个可以用点分治解决。
具体地，对于一个分治中心$root$，我们求出它到该分治结构内部所有点的路径权值（包含$root$），记为$diss_i$，和内部所有点到它的路径权值（不包含$root$），记作$dist_i$。那么一条路径$(u,v),\mbox{lca}(u,v)=root$的权值为

$$dist_u+diss_v\cdot k^{dep_u}$$
那么它模意义下等于 $x$就是
$$\frac{x-diss_u}{k^{dep_u}}=diss_v$$
开两个哈希表存式子分别左右两边的东西，每处理一棵子树先把其内部信息在哈希表中删去，然后在相互查询即可。
代码：

include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#define N 100010
#define ll long long
using namespace std;
const int R=5000000;
int tote,n,mod,K,X,val[N],to[N<<1],nxt[N<<1],con[N],sz[N],dfn[N],ed[N],fd[N],tim,cnts[N],cntt[N],dep[N];
ll diss[N],dist[N],pk[N],ik[N];
bool vis[N];
int read()
{
    int x=0,f=1;char ch=getchar();
    for(;ch<'0'||ch>'9';ch=getchar()) if(ch=='-') f=-1;
    for(;ch>='0'&&ch<='9';ch=getchar()) x=x*10+ch-'0';
    return x*f;
}
void ins(int x,int y)
{
    to[++tote]=y;
    nxt[tote]=con[x];
    con[x]=tote;
}
ll ksm(ll a,ll b)
{
    ll r=1;
    for(;b;b>>=1,a=a*a%mod)
        if(b&1) r=r*a%mod;
    return r;
}
struct ha
{

    int cnt,top,us[N],hd[R+10],f[N],s[N],nxt[N];
    ha(){cnt=top=0;}
    void add(int x,int d)
    {
        int id=x%R;
        for(int p=hd[id];p;p=nxt[p])
            if(f[p]==x) {s[p]+=d;return ;}
        if(hd[id]==0) us[++top]=id;
        s[++cnt]=d;f[cnt]=x;
        nxt[cnt]=hd[id];
        hd[id]=cnt; 
    }
    int qry(int x)
    {
        int id=x%R;
        for(int p=hd[id];p;p=nxt[p])
            if(f[p]==x) return s[p];
        return 0;   
    }
    void clr()
    {
        for(int i=1;i<=cnt;i++)
            f[i]=s[i]=nxt[i]=0;
        for(int i=1;i<=top;i++)
            hd[us[i]]=0,us[i]=0;
        top=cnt=0;                  
    }
}Hs,Ht;
void getroot(int v,int fa,int size,int &root)
{
    sz[v]=1;
    bool flag=1;
    for(int p=con[v];p;p=nxt[p])
        if(!vis[to[p]]&&to[p]!=fa)
        {
            getroot(to[p],v,size,root);
            sz[v]+=sz[to[p]];
            if((sz[to[p]]<<1)>size) flag=0;
        }
    if((sz[v]<<1)<size) flag=0;
    if(flag) root=v;    
}
void dfs(int v,int fa)
{
    dfn[v]=++tim;fd[tim]=v;
    diss[v]=(diss[fa]+pk[dep[v]]*val[v])%mod;
    dist[v]=fa?(dist[fa]*K+val[v])%mod:0;
    for(int p=con[v];p;p=nxt[p])
        if(!vis[to[p]]&&to[p]!=fa)
            dep[to[p]]=dep[v]+1,dfs(to[p],v);
    ed[v]=tim;
}
ll G(int v)
{
    return ik[dep[v]]*(X-dist[v]+mod)%mod;
}
void solve(int v,int size)
{
    int root=0;
    getroot(v,0,size,root);
    vis[root]=1;dep[root]=0;tim=0;
    dfs(root,0);
    for(int i=1;i<=size;i++)
        Hs.add(diss[fd[i]],1),Ht.add(G(fd[i]),1);
    cnts[root]+=Hs.qry(X),cntt[root]+=Ht.qry(val[root]);
    for(int p=con[root];p;p=nxt[p])
        if(!vis[to[p]])
        {
            int u=to[p];
            for(int i=dfn[u];i<=ed[u];i++)
                Hs.add(diss[fd[i]],-1),Ht.add(G(fd[i]),-1);
            for(int i=dfn[u];i<=ed[u];i++)
                cnts[fd[i]]+=Hs.qry(G(fd[i])),cntt[fd[i]]+=Ht.qry(diss[fd[i]]);
            for(int i=dfn[u];i<=ed[u];i++)
                Hs.add(diss[fd[i]],1),Ht.add(G(fd[i]),1);       
        }   
    Hs.clr();Ht.clr();
    for(int p=con[root];p;p=nxt[p])
        if(!vis[to[p]])
            solve(to[p],sz[to[p]]>sz[root]?size-sz[root]:sz[to[p]]);    
}
int main()
{
    freopen("el.in","r",stdin);
    freopen("el.out","w",stdout);
    n=read();mod=read();K=read();X=read();
    for(int i=1;i<=n;i++)
        val[i]=read();
    for(int i=1;i<n;i++)
    { 
        int x=read(),y=read();
        ins(x,y);ins(y,x);
    }   
    pk[0]=1;
    for(int i=1;i<=n;i++)
        pk[i]=pk[i-1]*K%mod;
    ik[n]=ksm(pk[n],mod-2);
    for(int i=n-1;i>=0;i--)
        ik[i]=ik[i+1]*K%mod;
    solve(1,n);
    ll ans=0;
    for(int i=1;i<=n;i++)
    {
        ans+=2ll*cnts[i]*(n-cnts[i]);
        ans+=2ll*cntt[i]*(n-cntt[i]);
        ans+=(ll)cnts[i]*(n-cntt[i]);
        ans+=(ll)cntt[i]*(n-cnts[i]);
    }   
    ans>>=1;ans=(ll)n*n*n-ans;
    printf("%lld",ans);
    return 0;
}