yukicoder : No.2514 Twelvefold Way Returns

これは知らない知識だった…。
https://yukicoder.me/problems/no/2514

問題

N個の区別できるボールをM個の箱に入れる。
各箱のボールの個数を3で割ると1余るような入れ方は何通りか。

解法

(x+x^4/4!+x^7/7!+...)^Mにおけるx^Nの係数を取れればよいのだが、いかんせんNが大きすぎてこのままでは計算できない。
1の3乗根ωを使うと(x+x^4/4!+x^7/7!+...)の部分が(e^x+ω^2*e^(ωx)+ω*e^(ω^2x))/3と置き換えられることを用いて式変形すると、(a+bω+cω^2)^Nの整数項を求める形にできる。

const ll mo=998244353;
int N,M;

const int NUM_=400001;
static ll fact[NUM_+1],factr[NUM_+1],inv[NUM_+1];

vector<ll> mult(vector<ll> A,vector<ll> B) {
	vector<ll> C(3);
	int a,b;
	FOR(a,3) FOR(b,3) (C[(a+b)%3]+=A[a]*B[b])%=mo;
	return C;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	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;
	
	cin>>N>>M;
	ll ret=0;
	int a,b,c;
	for(a=0;a<=M;a++) for(b=0;a+b<=M;b++) {
		c=M-a-b;
		vector<ll> A={a,b,c},B={(2*b+c)%3==0,(2*b+c)%3==1,(2*b+c)%3==2};
		x=N;
		while(x) {
			if(x%2) B=mult(A,B);
			A=mult(A,A);
			x/=2;
		}
		ll v=(mo+B[0]-B[1])*fact[M]%mo*factr[a]%mo*factr[b]%mo*factr[c]%mo;
		(ret+=v)%=mo;
	}
	FOR(i,M) ret=ret*(mo+1)/3%mo;
	cout<<ret<<endl;
	
	
}

まとめ

この式変形パターンを初めて見たので、こりゃ自力では解けないな…。