Official

Ex - add 1 Editorial by Nyaan


この解法は準備者の pencil さんの解説とおおむね一致しています。(実はこの問題は自分の原案なので、原案者としての想定解を簡単に説明します)

現在のカウンターの値を \(C_i\) とします。 \(D_i = C_i - A_i\) とします。列 \(D = (D_1, D_2, \dots, D_N)\) に注目すると、問題は次のように言い換えられます。

はじめ \(D = (A_1, \dots, A_N)\) である。次の一連の操作を \(\max(D) = 0\) になるまで行う。操作回数の期待値は?

  • \(D\) のすべての要素から \(1\) を引く。その後、\(1 \leq r \leq N\) である 整数 \(r\) をランダムに選び \(D_r \gets A_r\) とする。

\(\max(D)\) に注目します。操作後に max たりうる値は ((操作前に max だった値) \(- 1\)) または (\(r\) として選ばれた index に対応する \(D_r\) ) です。また、操作終了条件は \(\max(D) = 0\) です。よって \(\max(D)\) のみを状態として管理しても問題ないことが分かります。

このような言い換えを経ると、言い換え後の問題は吸収マルコフ連鎖で使えるテクニックを straightforward に適用できる形になります。
ポテンシャル関数 \(f\) を次のように置きます。(ポテンシャル関数の説明は ABC249 Ex : Dye Color の解説 などを参考にしてください。簡単に言うと、\(f(x)\) は (\(\max(D)=x\) から \(\max(D) = 0\) になるまでにかかる回数の期待値) に定数を足したもの、と考えて良いです。よって適切に定数を足せばよく見る期待値 DP の解法に一致しますが、余計な式変形が増えて面倒なのでここでは触れません。)

\[f(x) = 1 + \frac{1}{N} \sum_{i=1}^N f(\max(x-1,A_i)) \ (0 \leq x \lt A_N)\]

このとき \(f(A_N) - f(0)\) が答えです。よって \(f(A_N) = 0\) とおいたときの \(-f(0)\) を計算できればよいです。
\(f(A_N)\) から \(f(A_{N-1})\) を求める方法を説明します。(これと同様の計算を \(N-1\) 回繰り返せば \(f(A_1) = f(0)\) が求まります。)

\(A_{N-1} \lt x \leq A_N\) の場合

\[ \begin{aligned} &f(x) = 1 + \frac{N-1}{N} f(x-1) + \frac{1}{N} f(A_N) \\ &\iff f(x-1) = \frac{N}{N-1} \left(f(x) - 1 - \frac{f(A_N)}{N} \right) \end{aligned} \]

が成り立ちます。よって

\[a = \frac{N}{N-1}, b = -a\left(1+\frac{f(A_N)}{N}\right), G(x) = ax + b\]

とおくと

\[f(x-1) = G(f(x))\]

が成り立つので、\(G\)\(n\) 回合成した関数を \(G_n\) と表すと、 \(f(A_{N-1})\)\(f(A)\) の間には

\[ f(A_{N-1}) = G_{A_N - A_{N-1}}(A_N)\]

という関係式が成り立ちます。\(G_n\) は(陽に計算せずとも)行列累乗の要領で \(n\) に関する対数時間で計算できるので、\(f(A_{N-1})\)\(\mathrm{O}(\log (A_N - A_{N-1}))\) で求まります。これを繰り返せば \(f(0)\) が求まります。
以上よりこの問題を \(\mathrm{O}(N \log A_N)\) 程度で解くことができました。

  • 実装例
#include <iostream>
#include <vector>
using namespace std;
#include "atcoder/modint.hpp"
using mint = atcoder::modint998244353;

struct affine { mint a, b; };
affine operator*(affine l, affine r) { return {l.a * r.a, l.a * r.b + l.b}; }
affine affine_pow(affine a, long long e) {
  if (e == 0) return {1, 0};
  auto half = affine_pow(a, e / 2);
  return half * half * (e % 2 ? a : affine{1, 0});
}

int main() {
  int N;
  cin >> N;
  vector<long long> A(N);
  for (auto& x : A) cin >> x;
  mint s = 0, f = 0;
  for (int i = N - 1; i; i--) {
    mint a = mint{N} / i;
    mint b = -a * (s / N + 1);
    auto af = affine_pow({a, b}, A[i] - A[i - 1]);
    f = af.a * f + af.b, s += f;
  }
  cout << (-f).val() << endl;
}

posted:
last update: