kmjp's blog

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

AtCoder ABC #270 (トヨタ自動車プログラミングコンテスト2022) : Ex - add 1

解き切れてよかったね。
https://atcoder.jp/contests/abc270/tasks/abc270_h

問題

N要素の非負整数列Aが与えられる。
Aは昇順で、かつ(1-indexedで)A[1]=0である。

このAに対し以下を行う。

  • N要素で初期状態全部0のカウンタを考える。
  • 操作1回ごとに、カウンタを等確率で1つ選んで値を0にするとともに、他のカウンタを全部インクリメントする。
  • i個目のカウンタがA[i]以上であるような状態が全部のiで達成した段階で、操作を終える。

操作を完了するまでに行う操作回数の期待値を求めよ。

解法

自分の解法はEditorialの2つ目と同じ。
まずAを降順にしたうえでランレングス圧縮した数列(B[i],C[i])を考える。この数列の要素数をMとする。
これは、Aを降順で見て(重複を除き、0-indexedで)i番目となるものがB[i]であり、同じ値がC[i]個あることを示す。

状態遷移を考え、次の状態に移るための条件を考える。
i番目の状態から(i+1)番目に移るまでの処理回数の期待値f(i)を以下のように定義する。
f(i) := ここから、sum(C[i..(M-1)])個のカウンタを(B[i]-B[i+1])回連続で選べば、次の状態に進める。その時の処理回数の期待値

解はf(0)+f(1)+f(2)+....+f(M-1)となる。
i番目の状態から、誤ったカウンタを選んでしまうと、0~(i-1)の状態に戻ることがある。
よって、f(0)~f(i-1)の値に対し累積和や等比数列の総和の式を使いf(i)を求めよう。

int N;
ll A[202020];
const ll mo=998244353;

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 dp[202020];

void solve() {
	int i,j,k,l,r,x,y; 
	
	cin>>N;
	vector<pair<ll,int>> V;
	FOR(i,N) {
		cin>>A[i];
		if(V.empty()||V.back().first!=A[i]) V.push_back({A[i],0});
		V.back().second++;
	}
	reverse(ALL(V));
	
	int ok=N;
	ll penalty=0;
	ll ret=0;
	int s=0;
	FOR(i,V.size()-1) {
		ok-=V[i].second;
		ll turn=V[i].first-V[i+1].first;
		ll okp=ok*modpow(N)%mo;
		// sum((i+1)*ok^i)
		ll a=(1-modpow(okp,turn)+mo)*modpow(1-okp+mo)%mo+mo-turn*modpow(okp,turn)%mo+mo;
		a=a%mo*modpow(1-okp+mo)%mo;
		// sum(ok^i)
		ll b=(1-modpow(okp,turn)+mo)*modpow(1-okp+mo)%mo;
		
		(a*=(1-okp+mo))%=mo;
		(b*=(1-okp+mo))%=mo;
		ll C=turn*modpow(okp,turn)%mo;
		C+=a;
		C+=penalty*b%mo*modpow(N-ok)%mo;
		C=C%mo;
		dp[i]=C*modpow(1-b+mo)%mo;
		
		(penalty+=(s+V[i].second)*dp[i])%=mo;
		s+=V[i].second;
		ret+=dp[i];
	}
	cout<<ret%mo<<endl;
	
}

まとめ

解法が思いついた後、式変形にかなり手間取った。
実装ミスも多かったし、modint使えばよかったかな…。