[牛客网Wannafly挑战赛23F]计数

Description

给定一个n个点m条边的带权无向图，求边权和是k的倍数的生成树个数。
答案对p取模。
n,k<=100,m<=10000,p<=1e9且p是质数，p≡1(mod k)

Solution

套路计数题
考虑把一条边权为c的边看做单项式x^c，然后做矩阵树定理，答案就是所有x^nk的系数之和。
但是这个多项式的次数界可能有nk，直接插值的复杂度无法接受
考虑写成循环卷积的形式，那么答案就是常数项，问题变成了如何把普通卷积变成循环卷积
注意到我们最后多项式的次数界是k-1，也就是说插值需要k个点值
和NTT的方法相同，我们只需要选择k个满足x^k=1的x代入
而p是质数一定存在原根g，p-1又是k的倍数，这k个数显然就是g^(i(p-1)/k)，i=0~k-1了
最后拉格朗日插值求出常数项系数

Code

#include <cstdio>
#include <cstring>
#include <algorithm>
#define fo(i,a,b) for(int i=a;i<=b;i++)
#define fd(i,a,b) for(int i=a;i>=b;i--)
using namespace std;

typedef long long ll;

const int N=105,M=N*N;

int n,m,k,p,pri[N],u[M],v[M],w[M];
int a[N][N],G[N],an[N];

int pwr(int x,int y) {
    int z=1;
    for(;y;y>>=1,x=(ll)x*x%p)
        if (y&1) z=(ll)z*x%p;
    return z;
}

int get_g(int p) {
    int x=p,tot=0;
    for(int i=2;i*i<=x;i++) {
        if (!(x%i)) pri[++tot]=i;
        while (!(x%i)) x/=i;
    }
    if (x>1) pri[++tot]=x;
    for(int g=2;;g++) {
        bool ok=1;
        fo(i,1,tot) if (pwr(g,p/pri[i])==1) {ok=0;break;}
        if (ok) return g;
    }
}

int gauss(int n) {
    int sgn=1;
    fo(i,1,n) {
        if (!a[i][i])
            fo(j,i,n) 
                if (a[j][i]) {
                    fo(k,i,n) swap(a[j][k],a[i][k]);
                    sgn=-sgn;break;
                }
        int inv=pwr(a[i][i],p-2);
        fo(j,i+1,n) {
            int w=(ll)a[j][i]*inv%p;
            fo(k,i,n) (a[j][k]-=(ll)w*a[i][k]%p)%=p;
        }
    }
    int ret=1;
    fo(i,1,n) ret=(ll)ret*a[i][i]%p;
    return ret*sgn;
}

int main() {
    scanf("%d%d%d%d",&n,&m,&k,&p);
    fo(i,1,m) scanf("%d%d%d",&u[i],&v[i],&w[i]);
    int g=get_g(p-1);
    fo(i,0,k-1) {
        G[i]=pwr(g,(p-1)/k*i);
        memset(a,0,sizeof(a));
        fo(j,1,m) {
            int x=u[j],y=v[j],val=pwr(G[i],w[j]);
            (a[x][x]+=val)%=p;(a[y][y]+=val)%=p;
            (a[x][y]-=val)%=p;(a[y][x]-=val)%=p;
        }
        an[i]=gauss(n-1);
    }
    int ans=0;
    fo(i,0,k-1) {
        int ret=1;
        fo(j,0,k-1) if (j!=i) ret=(ll)ret*G[j]%p*pwr(G[j]-G[i],p-2)%p;
        (ans+=(ll)ret*an[i]%p)%=p;
    }
    printf("%d\n",(ans+p)%p);
    return 0;
}