kmjp's blog

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

AtCoder ABC #208 : F - Cumulative Sum

作ったライブラリが役に立った。
https://atcoder.jp/contests/abc208/tasks/abc208_f

問題

関数F(n,m)は以下のように定義される。

  • n=0のとき0
  • m=0の時n^K
  • nもmも正の時、F(n,m)=F(n-1,m)+F(n,m-1)

N,M,Kが与えられるので、F(N,M)を求めよ。

解法

F(N,0)はNのK次式である。
F(N,M)はF(N,M-1)の累積和を取ったものになるので、K+M次式になる。
そこで、F(N,M)について、N=0~(K+M)の(K+M+1)要素を求めてラグランジュ補間しよう。

ll N;
int M;
int K;
ll F[2525250][31];
const ll mo=1000000007;

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

ll lagrange(vector<ll>& P,ll x) {
	const int NUM_=3000003;
	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;
	}

	int i;
	int N=P.size();
	if(0<=x&&x<N) return P[x];
	vector<ll> R={1},F={1};
	for(i=N-1;i>=1;i--) R.push_back(R.back()*((x-i)%mo)%mo);
	
	ll p=1;
	ll ret=0;
	FOR(i,N) {
		ll a=p*R.back()%mo*factr[i]%mo;
		if((N-1-i)%2==0) a=a*factr[N-1-i]%mo;
		else a=a*(mo-factr[N-1-i])%mo;
		(ret+=a*P[i])%=mo;
		R.pop_back();
		(p*=(x-i)%mo)%=mo;
	}
	return ret%mo;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M>>K;
	for(i=0;i<=K+M+1;i++) F[i][0]=modpow(i,K);
	for(i=1;i<=M;i++) {
		for(j=1;j<=K+M+1;j++) {
			F[j][i]=F[j-1][i]+F[j][i-1];
			if(F[j][i]>=mo) F[j][i]-=mo;
		}
	}
	
	vector<ll> P;
	FOR(i,K+M+2) P.push_back(F[i][M]);
	cout<<lagrange(P,N)<<endl;
	
}

まとめ

Kが微妙に大きいの何なんだろうな。
10^6位でもよい気がするが。