kmjp's blog

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

Codeforces ECR #007 : F. The Sum of the k-th Powers

日本人が妙に上位に並んだの、宜なるかな。
http://codeforces.com/contest/622/problem/F

問題

正整数N,Kに対し、 \displaystyle \sum_{i=1}^N i^Kを答えよ。

解法

N≦Kの時は愚直に計算すればよい。

それ以外の場合、 \displaystyle  f(x) = \sum_{i=1}^x i^Kと置くと、この関数は恐らく(K+1)次の多項式になることが推測がつく。
ということは、f(0)~f(K+1)の(K+2)個の値を求めれば、ラグランジュ補間でf(N)を求められる。

p次式に対し、f(0)~f(p)の値からO(p)でラグランジュ補間を行うテクは以下既出である。
AtCoder ARC #033 : D - 見たことのない多項式 - kmjp's blog

最初書いた解が手元で2.3s(CFサーバでは1.5s)だったので、下記の解は無駄にインラインアセンブラで除算を高速化している。
なくてもCFサーバなら通る。

ll N,K;
ll mo=1000000007;
ll V[1010101];
ll fact[1010101];

inline int mulmod(int a,int b,int mo) {
	int d,r;
	if(a==0 || b==0) return 0;
	if(a==1 || b==1) return max(a,b);
	__asm__("mull %4;"
	        "divl %2"
		: "=d" (r), "=a" (d)
		: "r" (mo), "a" (a), "d" (b));
	return r;
}

int modpow(int a, int n = mo-2) {
	int r=1;
	while(n) r=mulmod(r,(n%2)?a:1,mo),a=mulmod(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+3;i++) {
		V[i]=(V[i-1]+modpow(i,K))%mo;
		fact[i]=(fact[i-1]*i)%mo;
	}
	
	if(N<=K+2) {
		cout<<V[N]<<endl;
		return;
	}
	
	K++;
	ll a=1;
	FOR(i,K+1) a=a*(N+mo-i)%mo;
	ll ret=0;
	
	FOR(i,K+1) {
		int b=mulmod(fact[i],fact[K-i],mo);
		ll rev2=N-i;
		if((K-i)%2) b=mo-b;
		ret += mulmod(mulmod(mulmod(V[i],a,mo),modpow(rev2),mo),modpow(b),mo);
	}
	cout<<ret%mo<<endl;	
	
}

まとめ

気づけて良かった。