解き切れてよかったね。
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使えばよかったかな…。