kmjp's blog

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

yukicoder : No.665 Bernoulli Bernoulli

そそっかしい+数学弱者+記憶力不足と色々やらかした。
https://yukicoder.me/problems/no/665

問題

正整数n,kが与えられる。
 \displaystyle \sum_{i=1}^n i^k \mod 1000000007を答えよ。

解法

問題をしっかり読むと、問題名からベルヌーイ数がヒントであることがわかる。

自分はそこに気付かず、そもそもベルヌーイ数とこのような総和の関係をちゃんと理解していないので本番ラグランジュ補間で解いた。
 \displaystyle S(x) = \sum_{i=1}^x i^k \mod 1000000007とすると、S(x)は(k+1)次式になることが想像できるのでS(0)~S(k+1)を求めればラグランジュ補間でS(N)も求められる。

愚直に計算するとO(k^2)回剰余計算を行うのでTLEする。
ラグランジュ補間の過程では似たような数値や階乗の乗算を繰り返すので、うまく前処理するとO(k)に押さえることができる。

本番は以下の問題を思い出しつつ計算量を頑張って落として通した。
AtCoder ARC #033 : D - 見たことのない多項式 - kmjp's blog

残念ながら計算量を落とすテクは覚えていたのに、この問題の詳細は覚えていなかった。
この問題はまさにラグランジュ補間そのものの問題なので、S(0)~S(k+1)を求めた後は上記問題コードを流用すれば良かったのだ…。

さらにその後、まったく同じでむしろ制限の厳しい問題が既出であることを知った。
Codeforces ECR #007: F. The Sum of the k-th Powers - kmjp's blog

ll N,K;
ll mo=1000000007;
ll P[101010];
ll fact[101010];

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

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	
	fact[0]=1;
	for(i=1;i<=K+1;i++) {
		(P[i]=modpow(i,K)+P[i-1])%=mo;
		fact[i]=fact[i-1]*i%mo;
	}
	
	if(N<=K+1) return _P("%lld\n",P[N]);
	N=N%mo+mo;
	
	ll ret=0;
	ll A=1;
	for(i=0;i<=K+1;i++) A=A*(N-i)%mo;
	
	for(i=0;i<=K+1;i++) {
		ll v=P[i]*A%mo*modpow(N-i)%mo;
		ll w=fact[K+1-i]*fact[i]%mo;
		v=v*modpow(w)%mo;
		
		if(i%2 != (K+1)%2) ret=(ret+mo-v)%mo;
		else ret=(ret+v)%mo;
	}
	
	cout<<ret<<endl;
	
}

まとめ

書いた本人が覚えてないブログエントリを引っ張り出してくる人々、いったいなんなんだ…。