yukicoder : No.2833 Count Taiko Results

少し言い換えると楽。
https://yukicoder.me/problems/no/2833

問題

N個のノーツからなる音ゲーを考える。
i番目のノーツでは、A[i]+B[i]段階の判定があり、上位A[i]段は成功、下位B[i]段は失敗とする。
連続成功がちょうどK回であるようなゲームの展開パターンは何通りか。

解法

連続成功がK回以下であるケースから、(K-1)回以下であるケースを引こう。

あらかじめA[i]・B[i]をA[i]で割っておく、すなわち各ノーツでの成功率を1、失敗率をB[i]/A[i]としておくと計算しやすい。
dp(n,k) := 連続正解がk回以下であるとき、n回目までのノーツのパターンを定めたときにn回目が失敗であるような確率
とし、累積和を使ってdp(n-k,k)～dp(n,k)からdp(n+1,k)を求めて行けばよい。

int N,K;
int A[202020],B[202020];
const ll mo=998244353;
ll dp[202020];
ll dp2[202020];

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;
	
	dp[0]=dp2[0]=1;
	cin>>N>>K;
	FOR(i,N) cin>>A[i];
	FOR(i,N) cin>>B[i];
	ll ret=0;
	
	if(K==N) ret=1;
	
	
	FOR(i,N) {
		j=B[i]*modpow(A[i])%mo;
		x=(mo+dp[i]-((i-(K+1)>=0)?dp[i-(K+1)]:0))*j%mo;
		y=(mo+dp2[i]-((i-K>=0)?dp2[i-K]:0))*j%mo;
		dp[i+1]=(dp[i]+x)%mo;
		dp2[i+1]=(dp2[i]+y)%mo;
		if(i+K>=N-1) ret+=x;
		if(i+K-1>=N-1) ret+=mo-y;
	}
	ret%=mo;
	FOR(i,N) ret=ret*A[i]%mo;
	cout<<ret<<endl;
	
	
}

まとめ

先にA[i]で割らないと、ちょっと面倒くさそう。