牛客网Wannafly挑战赛23 F-计数 矩阵树定理+拉格朗日插值法

版权声明：本文为博主原创文章，未经博主允许不得转载。 https://blog.csdn.net/qq_33229466/article/details/82559329

题意

给出一个n个点m条边的带权无向图，问有多少棵生成树满足边权和模k等于0，答案模p输出。
$n,k\le100,p\le10^9,p为质数且满足k\equiv1\pmod p$。

分析

一个显然的做法就是把每条边的边权看成是$x^w$，然后做矩阵树定理，最后得出来多项式的$x^{ik}$项系数的和就是答案。
但答案多项式的次数是$nk$级别的，显然不能做。注意到如果我们把卷积变成模$x^k$意义下的循环卷积，那么最后常数项即为答案。
显然最后多项式的次数是$k-1$，我们可以考虑求出$k$个点值然后插值。
问题在于如何模拟循环卷积。如果我们像NTT那样，选择$k$个不同的满足$x^k\equiv 1\pmod p$的$x$作为横坐标，然后带进去就可以得到循环卷积后的点值了。
注意到$p$是质数所以必然存在原根，又有$p-1$是$k$的倍数，所以$x$就可以取$g^{i(p-1)/k}$，其中$i=0,1,...,k-1$。
时间复杂度$O(n^3k)$。

代码

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

typedef long long LL;

const int N=105;

int n,m,k,MOD,stack[25],xi[N],yi[N],a[N][N];
struct edge{int x,y,w;}e[N*N];

int read()
{
    int x=0,f=1;char ch=getchar();
    while (ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
    while (ch>='0'&&ch<='9'){x=x*10+ch-'0';ch=getchar();}
    return x*f;
}

int ksm(int x,int y)
{
    int ans=1;
    while (y)
    {
        if (y&1) ans=(LL)ans*x%MOD;
        x=(LL)x*x%MOD;y>>=1;
    }
    return ans;
}

int get_g(int n)
{
    int tmp=n-1,top=0;
    for (int i=2;i*i<=tmp;i++)
        if (tmp%i==0)
        {
            stack[++top]=i;
            while (tmp%i==0) tmp/=i;
        }
    if (tmp>1) stack[++top]=tmp;
    for (int i=2;;i++)
    {
        bool flag=0;
        for (int j=1;j<=top;j++)
            if (ksm(i,(n-1)/stack[j])==1) {flag=1;break;}
        if (!flag) return i;
    }
}

int gauss(int n)
{
    int ans=1;
    for (int i=1;i<=n;i++)
    {
        if (!a[i][i])
        {
            for (int j=i+1;j<=n;j++)
                if (a[j][i])
                {
                    for (int k=i;k<=n;k++) std::swap(a[j][k],a[i][k]);
                    ans=-ans;
                    break;
                }
        }
        for (int j=i+1;j<=n;j++)
            if (a[j][i])
            {
                int w=(LL)a[j][i]*ksm(a[i][i],MOD-2)%MOD;
                for (int k=i;k<=n;k++) (a[j][k]+=MOD-(LL)a[i][k]*w%MOD)%=MOD;
            }
        ans=(LL)ans*a[i][i]%MOD;
    }
    return (ans+MOD)%MOD;
}

int main()
{
    n=read();m=read();k=read();MOD=read();
    int g=get_g(MOD);
    for (int i=1;i<=m;i++) e[i].x=read(),e[i].y=read(),e[i].w=read();
    for (int i=1;i<=k;i++)
    {
        xi[i]=ksm(g,(MOD-1)/k*(i-1));
        memset(a,0,sizeof(a));
        for (int j=1;j<=m;j++)
        {
            int u=e[j].x,v=e[j].y,w=e[j].w,t=ksm(xi[i],w);
            (a[u][u]+=t)%=MOD;(a[v][v]+=t)%=MOD;
            (a[u][v]+=MOD-t)%=MOD;(a[v][u]+=MOD-t)%=MOD;
        }
        yi[i]=gauss(n-1);
    }
    int ans=0;
    for (int i=1;i<=k;i++)
    {
        int w=yi[i];
        for (int j=1;j<=k;j++)
            if (i!=j) w=(LL)w*(MOD-xi[j])%MOD*ksm(xi[i]+MOD-xi[j],MOD-2)%MOD;
        (ans+=w)%=MOD;
    }
    printf("%d\n",ans);
    return 0;
}