yukicoder : No.2530 Yellow Cards

今回はサクサク解ける回だったね。
https://yukicoder.me/problems/no/2530

問題

正整数N,Kが与えられる。
初期状態で1～N番の背番号を持つN人の選手がプレイエリアにいる。

その後、プレイエリアにいる選手のうち誰か1人、等確率でイエローカードが提示される。
イエローカードが2枚提示された選手はプレイエリアから退場し、代わりに(N+1)番以降で未出の背番号の選手がプレイエリアに入る。
K回イエローカードが提示されたあと、プレイエリアにいる人の背番号の最大値の期待値を求めよ。

解法

f(n,k) := k回イエローカードが提示された状態で、1枚イエローカードが提示された人がn人いるような状態に至る確率
を考える。
退場済みの人数は(k-n)/2となるし、その時プレイエリアにいる選手の背番号の最大値はN+(k-n)/2である。
よって解はである。
あとはf(n,k)を埋めて行こう。

int N,K;
const ll mo=998244353;

ll from[5050];
ll to[5050];

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;
	from[0]=1;
	ll re=modpow(N);
	FOR(i,K) {
		ZERO(to);
		FOR(j,K+1) if(from[j]) {
			if(j) (to[j-1]+=from[j]*j%mo*re)%=mo;
			if(j<N) (to[j+1]+=from[j]*(N-j)%mo*re)%=mo;
		}
		
		swap(from,to);
	}
	ll ret=0;
	FOR(i,5010) if(from[i]) (ret+=from[i]*((K-i)/2+N))%=mo;
	cout<<ret<<endl;
	
}

まとめ

こちらは割とすんなり。