kmjp's blog

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

CSAcademy Round #68 : E. Sliding Product Sum

勉強になりました。
https://csacademy.com/contest/round-68/task/sliding-product-sum/

問題

整数N,K,Mが与えられる。
[1,2,...,N]の数列に対し、このうち長さK以下の連続した部分列に考え、その積の総和をmod Mで答えよ。

解法

長さに関し1~Kそれぞれのケースを考え加算しよう。
以下長さがちょうどKのケースを考える。

長さがKである連続部分列の和は、 \displaystyle {}_K P_K + {}_{K+1} P_K + ... + {}_N P_K = K! \times \left( {}_K C_0 + {}_{K+1} C_1 + ... + {}_{K+(N-K)} K_{N-K} \right)となる。
この括弧内をどうにかしよう。
実はこの括弧内は {}_{K+(N-K)+1 K_{N-K} = {}_{N+1} K_{K+1}と置ける。
これはパスカルの三角形を眺めても推測が付くし、先頭の {}_K C_0 {}_{K+1} C_0と読み替えると、先頭の2項をパスカルの三角形の要領で足しこんでいくと結局全部が1個の項になる。

この性質は自分は記載を見つけられなかったが、Editorialによると英語Wikipediaに記載があるようだ。
Hockey-stick identity - Wikipedia

こうすると元の式は \displaystyle {}_K P_K + {}_{K+1} P_K + ... + {}_N P_K = K! \times {}_{N+1} K_{K+1} = \frac{(N+1)!}{(N-K)!(K+1)}となるのでそれを計算すればよい。
分子は連続する(K+1)要素の積なので、どこか1個は(K+1)で割れる値が登場する。

ll N,K,M;

ll comb(ll a,ll b) {
	ll bb=b;
	__int128_t r=1;
	while(b--) {
		if(bb>1 && a%bb==0) {
			r=r*(a/bb)%M;
			bb=1;
		}
		else {
			r=r*a%M;
		}
		a--;
	}
	return r;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K>>M;
	ll ret=0;
	for(i=1;i<=K;i++) (ret+=comb(N+1,i+1))%=M;
	cout<<ret<<endl;
	
}

まとめ

この二項係数の和の特性、過去にも使った気がしたものの思い出せなかったが、これで記憶に残りそう。