kmjp's blog

競技プログラミング参加記です

AtCoder ABC #279 (トヨタシステムズプログラミングコンテスト2022) : Ex - Sum of Prod of Min

ARC

また知らん知識が出てきた…。
https://atcoder.jp/contests/abc279/tasks/abc279_h

問題

整数N,Mが与えられる。MはN以上2N以下とする。
N要素からなる正整数列Sで、和がMとなるものを考える。
そのようなS全通りにおいて、Prod(min(i,S[i]))の総和を200003で割った余りを求めよ。

解法

数列の和に関するFPSで考えると、S[i]の寄与は
\displaystyle x+2x^2+...ix^i+ix^{i+1}+... = \frac{x(1-x^i)}{(1-x)^2}となる。
i=1～Nに対しこれらの積を取り、x^Mの係数を取ることを考えると、
\displaystyle [x^M]\prod_i \frac{x(1-x^i)}{(1-x)^2} = [x^{M-N}]\frac{\prod (1-x^i)}{(1-x)^{2N}}
ここでオイラーの五角数定理より、a_i = (-1)^i, e_i=i(3i-1)/2として上の式は
\displaystyle [x^{M-N}]\frac{\sum_{i=-\infty}^{\infty} a_i x^{e_i}}{(1-x)^{2N}}

iを、プラスマイナスともにe_iがM-N以下に収まる範囲で探索すると、高々O(√N)通り探索するとよい。
\displaystyle \frac{a_i x^{e_i}}{(1-x)^{2N}}のx^{M-N}への寄与は、 \displaystyle a_i \times Comb(M+N-e_i-1, N-1)である。
この変形は、ABC230ExのEditorialのMULTISETに関する解説でわかる。
Editorial - AtCoder Beginner Contest 230

なお、この二項係数の計算では、200003以上の値に対する二項係数を計算する必要がある。
そこはLucasの定理を用いて計算しよう。

ll N,M;
const ll mo=200003;

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;
}

ll hoge(ll P,ll Q) {
	ll ret=1;
	if(Q<0||Q>P) return 0;
	while(P) {
		(ret*=comb(P%mo,Q%mo))%=mo;
		P/=mo;
		Q/=mo;
	}
	return ret;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M;
	ll ret=0;
	for(ll k=0;;k++) {
		if(k) {
			ll e=k*(3*k+1)/2;
			if(e<=M-N) {
				ll a=hoge(M+N-e-1,2*N-1);
				if(k%2==1) {
					ret+=mo-a;
				}
				else {
					ret+=a;
				}
			}
		}
		{
			ll e=k*(3*k-1)/2;
			if(e>M-N) break;
			ll a=hoge(M+N-e-1,2*N-1);
			if(k%2==1) {
				ret+=mo-a;
			}
			else {
				ret+=a;
			}
		}

	}
	cout<<ret%mo<<endl;
}

解法

1/((1-x)^{2N})のところを二項係数に置き換えるところ、最初解説見てわからなくて、FPSを丁寧に読み取いたら重複組み合わせっぽい感じにおけそう…と思って調べたら、ちゃんと過去のABCの解説にあった。