AtCoder ABC #386 : G - Many MST

最終的なコードは割と短い。
https://atcoder.jp/contests/abc386/tasks/abc386_g

問題

正整数N,Mが与えられる。
N点の完全無向グラフで、各辺の重みが1～Mの整数のいずれかであるとき、あり得る全グラフにおける最小全域木の重みの総和を求めよ。

解法

重さm以下の辺だけ残したとき、ある頂点集合sが連結になってかつ残り(N-s)点と連結でないようなパターンdp(m,s)を数え上げる。
s点と(N-s)点の間は重さm+1以上でないといけないので、dp(m,s)*Comb(N,s)*(M-m)^(s*(N-s))*M^((N-s)*(N-s-1)/2)だけ解に寄与する。
dp(m,s)は、dp(m,1)～dp(m,s-1)を用いて計算できる。(1番の頂点が、1～(s-1)個の頂点までしか連結にできないケースを引く)

int N,M;
const ll mo=998244353;

ll F[505];
int po[502][505*260];
ll modpow(ll a, ll n = mo-2) {
	ll r=1;a%=mo;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

ll comb(ll N_, ll C_) {
	const int NUM_=400001;
	static ll fact[NUM_+1],factr[NUM_+1],inv[NUM_+1];
	if (fact[0]==0) {
		inv[1]=fact[0]=factr[0]=1;
		for (int i=2;i<=NUM_;++i) inv[i] = inv[mo % i] * (mo - mo / i) % mo;
		for (int i=1;i<=NUM_;++i) fact[i]=fact[i-1]*i%mo, factr[i]=factr[i-1]*inv[i]%mo;
	}
	if(C_<0 || C_>N_) return 0;
	return factr[C_]*fact[N_]%mo*factr[N_-C_]%mo;
}


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M;
	ll pat=(N-1-M+mo)*modpow(M,N*(N-1)/2)%mo;
	
	FOR(i,502) {
		po[i][0]=1;
		FOR(j,501*255) po[i][j+1]=1LL*po[i][j]*i%mo;
	}
	
	for(k=1;k<=M;k++) {
		for(int s=1;s<=N;s++) {
			F[s]=po[M][s*(s-1)/2];
			for(i=1;i<s;i++) F[s]-=F[i]*comb(s-1,i-1)%mo*po[M-k][i*(s-i)]%mo*po[M][(s-i)*(s-i-1)/2]%mo;
			F[s]=(F[s]%mo+mo)%mo;
			(pat+=F[s]*comb(N,s)%mo*po[M-k][s*(N-s)]%mo*po[M][(N-s)*(N-s-1)/2])%=mo;
		}
	}
	
	
	cout<<pat<<endl;
	
}

まとめ

これは落ち着いてやれば解けそうな問題だったな…。