問題

$n$ 人の人が一列に並んでおり、人 $i$ は先頭から $i$ 番目に並んでいる。以下の操作を人が1人になるまで行う。

先頭に並んでいる人を、確率 $\displaystyle \frac{1}{2}$ で列から取り除き、確率 $\displaystyle \frac{1}{2}$ で列の最後尾に移す。

人 $1, 2, \cdots, n$ について、人 $i$ が最後の1人になる確率を $\mathrm{mod} \ 998244353$ で求めよ。

補足

この問題では、求める確率が有理数であることが保証される。

また、確率を既約分数で $\displaystyle \frac{y}{x}$ と表したときに、 $x$ が $998244353$ で割り切れないことが保証される。

このときに、 $xz \equiv y \ (\mathrm{mod} \ 998244353)$ が成立する整数 $z$ （ただし、 $0 \leq z \lt 998244353$ ）が一意に定まるので、これを求めよ。

入力

最初の1行に、整数 $n$ が与えられる。

条件

実行時間制限: 2s
メモリ制限: 1024MB
$2 \leq n \leq 3000$

解法

まず、分数を $\mathrm{mod} \ 998244353$ で表すために、フェルマーの小定理 *1 を用います。

問題文に記載の通り、解答となる確率が $\displaystyle \frac{y}{x}$ のとき、 $x$ は $998244353$ （ $= m$ とおきます）と互いに素となります。従って、 ${x}^{m - 1} \equiv 1 \ (\mathrm{mod} \ m)$ が成立することから、便宜的に ${x}^{m - 2} \equiv {x}^{-1}$ とすることができます。

よって、求める値 $z$ については、 $z = ({x}^{m - 2} y) \% m$ ということになります。

さてここからは、この問題で求めるべき確率について考えます。

この問題では、「最後尾に移す」という操作を引き続ければ、何回でも操作を行うことができてしまいます。そのため、例えば「人 $i$ が最後の1人になるまでに行った操作が $k$ 回である確率」のようなものを求めていく方法では、プログラム上で正しい有理数の答えは得られません。

そこで、まずは確率 ${p}_{i}$ を、「 $i$ 人の人が並んでいるときに、人 $1$ が最後まで残る確率」として定義します。これを用いて、「 $n$ 人が並んでいるときに、人 $i$ が先頭になる確率」はどのように表せるか考えます。

まずは、人 $i$ が操作される直前に、人 $i$ 以前の人たちが脱落しているのか最後尾にいるのかについて考えます。

人 $i$ 以前にいる $i - 1$ 人にについて、このうち $j$ 人が列に残らず脱落したとすると、人 $i$ が先頭になったときには列に $(n - j)$ 人が残っていることになります。

従って、先程の $p$ を用いると、この状況になった際に、先頭の人 $i$ が最後まで列に残る確率は ${p}_{n - j}$ として表せます。

また、 $i - 1$ 人から $j$ 人が脱落する確率については、脱落する確率が $\displaystyle \frac{1}{2}$ であり、最後尾に移される確率が $\displaystyle \frac{1}{2}$ であることから、 \begin{align} \frac{(i - 1)!}{j! (i - 1 - j) !} \cdot {\left( \frac{1}{2} \right)}^{j} \cdot {\left( \frac{1}{2} \right)}^{i - 1 - j} \ = \ _{i - 1}\mathrm{C}_{j} {\left( \frac{1}{2} \right)}^{i-1} \end{align} として計算ができます。

以上から、「人 $i$ の最初の出番直前に $n - j$ 人が残っていて、そこから人 $i$ が最後まで残る確率」は、 \begin{align} _{i - 1}\mathrm{C}_{j} {\left( \frac{1}{2} \right)}^{i-1} {p}_{n - j} \end{align} となります。

この状況を図で整理すると以下のようになります。

図に示されている通り、 $j$ の値は $0 \leq j \leq i - 1$ の範囲であることから、 $n$ 人並んでいるときに人 $i$ が最後の1人になる確率は、 \begin{align} {\sum_{j = 0}^{i - 1}} \ & {_{i - 1}\mathrm{C}_{j}} {\left( \frac{1}{2} \right)}^{i-1} {p}_{n - j} \\ = & {_{i - 1}\mathrm{C}_{0}} {\left( \frac{1}{2} \right)}^{i-1} {p}_{n} + {_{i - 1}\mathrm{C}_{1}} {\left( \frac{1}{2} \right)}^{i-1} {p}_{n - 1} + \cdots + {_{i - 1}\mathrm{C}_{i - 1}} {\left( \frac{1}{2} \right)}^{i-1} {p}_{n - i + 1} \end{align} として計算することができます。

以上から、人 $i$ が最後の1人になる確率は、

$i = 1$ が最後の1人になる確率: ${p}_{n}$
$i = 2$ が最後の1人になる確率: $\displaystyle {_{1}\mathrm{C}_{0}} {\left( \frac{1}{2} \right)} {p}_{n} + {_{1}\mathrm{C}_{1}} {\left( \frac{1}{2} \right)} {p}_{n - 1}$
$\cdots \cdots$
$i = n$ が最後の1人になる確率: $\displaystyle {_{n - 1}\mathrm{C}_{0}} {\left( \frac{1}{2} \right)}^{n - 1} {p}_{n} + \cdots + {_{n - 1}\mathrm{C}_{n - 1}} {\left( \frac{1}{2} \right)}^{n - 1} {p}_{1}$

というように計算されます。いずれかの人が必ず最後の1人になることから、この値の総和は必ず $1$ となります。

ここで、便宜的に ${_{0}\mathrm{C}_{0}} = 1$ として以上のことを整理すると、 \begin{align} \sum_{i = 1}^{n} \sum_{j = 0}^{i - 1} {_{i - 1}\mathrm{C}_{j}} {\left( \frac{1}{2} \right)}^{i-1} {p}_{n - j} = 1 \end{align} が成立すると言えます。

この式の左辺に対して ${p}_{n - j}$ の係数に当たる部分を ${c}_{n, j}$ として定義します。すなわち、 \begin{align} \sum_{j = 0}^{n - 1} {c}_{n, j} \ {p}_{n - j} = 1 \end{align} となるような ${c}_{n, j}$ を定義すると、先程の整理から \begin{align} {c}_{n, j} = {_{j}\mathrm{C}_{j}} {\left( \frac{1}{2} \right)}^{j} + {_{j + 1}\mathrm{C}_{j}} {\left( \frac{1}{2} \right)}^{j + 1} + \cdots + {_{n - 1}\mathrm{C}_{j}} {\left( \frac{1}{2} \right)}^{n - 1}
\end{align} と言えます。

ここで行列の長さが $n + 1$ のときの、 ${p}_{n - j + 1}$ の係数に当たる部分である ${c}_{n + 1, j}$ の値について考えると、 \begin{align} {c}_{n + 1, j} & = {_{j}\mathrm{C}_{j}} {\left( \frac{1}{2} \right)}^{j} + {_{j + 1}\mathrm{C}_{j}} {\left( \frac{1}{2} \right)}^{j + 1} + \cdots + {_{n - 1}\mathrm{C}_{j}} {\left( \frac{1}{2} \right)}^{n - 1} + {_{n}\mathrm{C}_{j}} {\left( \frac{1}{2} \right)}^{n} \\ & = {c}_{n, j} + {_{n}\mathrm{C}_{j}} {\left( \frac{1}{2} \right)}^{n} \end{align} として整理できます。

また、 \begin{align} {c}_{n, n - 1} = {_{n - 1}\mathrm{C}_{n - 1}}{\left( \frac{1}{2} \right)}^{n - 1}
\end{align} であることから、 $j \leq n$ である $j$ に対して ${c}_{n, j} = 0$ として定めると、 \begin{align} {c}_{n + 1, n} = {_{n}\mathrm{C}_{n}}{\left( \frac{1}{2} \right)}^{n} = {c}_{n, n - 1} + {_{n}\mathrm{C}_{n}}{\left( \frac{1}{2} \right)}^{n} \end{align} と整理できるので、先ほどの漸化式が成立します。

ここで、行列の長さが $1$ のときについては、必ず人 $1$ が最後の1人になり、 ${p}_{1} = 1$ が成立することから、 ${c}_{1, 0} = 1$ となります。

初期値と漸化式があることから、 ${c}_{i, j}$ の値が求まることになります。

従って、 ${p}_{1}, {p}_{2}, \cdots , {p}_{i - 1}$ の値が全て求まっていると仮定すると、 \begin{align} \sum_{j = 0}^{i - 1} {c}_{i, j} \ {p}_{i - j} = 1 \end{align} より、 \begin{align} {c}_{i, 0} \ {p}_{i} = 1 - \sum_{j = 1}^{i - 1} {c}_{i, j} {p}_{i - j} \end{align} が成立するので、 ${p}_{i}$ が求まることになります。

よって、 ${c}_{1, 0} = 1, {p}_{1} = 1$ であることから、 $i = 2, 3, \cdots , n$ に対して ${c}_{i, j}$ の値を求め、それをもとに ${p}_{i}$ の値が求まることになります。

${_{i}\mathrm{C}_{j}}$ の値が既に求まっているとすると、 ${c}_{i, j}$ のそれぞれの値は $O(1)$ で求めることができます。その上で、 ${p}_{i}$ の値は $O(i)$ で求まるため、 ${p}_{n}$ の値が求まるまでに $O({n}^{2})$ の計算量が必要になります。

この ${p}_{1}, \cdots, {p}_{n}$ の値が求められると、先程の通り人 $i$ が最後になる確率は \begin{align} {\sum_{j = 0}^{i - 1}} \ {_{i - 1}\mathrm{C}_{j}} {\left( \frac{1}{2} \right)}^{i-1} {p}_{n - j} \end{align} によって求められるので、すべての人の確率についても $O({n}^{2})$ で求められます。

ここで、 ${_{n}\mathrm{C}_{k}}$ の値について、 \begin{align} {_{n}\mathrm{C}_{k}} = \frac{n \times \cdots \times (n - k + 1)}{i \times \cdots \times 1} \end{align} であるので、 \begin{align} {_{n}\mathrm{C}_{k + 1}} = \frac{n \times \cdots \times (n - k + 1) \times (n - k)}{(i + 1) \times i \times \cdots \times 1} = \frac{n - k}{i + 1} \cdot {_{n}\mathrm{C}_{k}} \end{align} と整理できます。これを用いることにより、それぞれの ${_{n}\mathrm{C}_{k}}$ の値を求めることができます。

以上から、大体 $O({n}^{2})$ の計算量ですべての確率を求めることができ、 $n \leq 3000$ なので、これは実行時間制限に間に合うと言えます。

ソースコード

まず、 ${x}^{y} \ (\mathrm{mod} \ 998244353)$ の値を高速に求める my_pow() 関数を以下のように実装しました。

const ll MOD2 = 998244353;

ll my_pow(ll x, ll y) {
  if (y == 0) {
    return 1; // 0乗は必ず 1 なので、1を返す
  }
  if (y % 2 == 0) {
    return my_pow(x * x % MOD2, y / 2); // y が偶数のときは (x ^ 2)^(y / 2) の値が答えとなる
  } else {
    return x * my_pow(x * x % MOD2, y / 2) % MOD2; // y が奇数のときは x * (x ^ 2)^(y / 2) の値が答えとなる
  }
}

この my_pow() 関数がある状態で、 $\displaystyle\frac{a}{b} \ (\mathrm{mod} \ 998244353)$ の値を求める frac() 関数を以下のように実装しました。

ll frac(ll a, ll b) {
  return a * my_pow(b, MOD2 - 2) % MOD2; // (1 / b) は mod m にて b ^ (m - 2) と同等であることを利用する
}

以上の関数が実装された状態で、答えを出力する部分を main() 関数の中に直接実装しました。

int n;
ll p[3005], ncr[3005][3005], coef[3005][3005]; // ncr は nCr の値を保存し、 coef は c[i][j] の値を保存する

int main() {
  // 予め 0 から 3000 までの nCr の値を計算する
  for (ll i = 0; i <= 3000; i++) {
    ncr[i][0] = 1; // nC0 = 1 であることを初期状態として持っておく
    for (ll j = 1; j <= i; j++) {
      ncr[i][j] = frac(i - j + 1, j) * ncr[i][j - 1] % MOD2; // 漸化式通りに値の更新をする
    }
  }

  // 1 から 3000 までの coef[i][j] の値を求めておく
  for (ll i = 1; i <= 3000; i++) {
    ll exp = my_pow(frac(1, 2), i - 1); // (1 / 2)^(i - 1) の値を予め計算しておく
    for (ll j = 0; j <= i - 1; j++) {
      coef[i][j] = (coef[i - 1][j] + ncr[i - 1][j] * exp) % MOD2; // 漸化式通りに値の更新をする
    }
  }

  // 1 から 3000 までの p[i] の値を順に計算する
  for (ll i = 1; i <= 3000; i++) {
    ll res = 1; // 方程式の右辺の値をもつ
    for (ll j = 1; j <= i - 1; j++) {
      res -= (p[i - j] * coef[i][j]) % MOD2; // 方程式通りに p[i - j] * coef[i][j] を引いていく
      if (res < 0) {
        res += MOD2; // 値が負になった場合は、 998244353 を足して0以上にする
      }
    }
    p[i] = frac(res, coef[i][0]); // p[i] の値を求める
  }

  cin >> n; // n の値を入力する

  // i = 1, ..., n の順に、それぞれ人 i が最後の1人になる確率を求める
  for (ll i = 1; i <= n; i++) {
    ll ans = 0; // 答えとなる値
    ll exp = my_pow(frac(1, 2), i - 1); // (1 / 2)^(i - 1) の値を求めておく

    // j = 0, ..., i - 1 の順に値を足していき、確率を求める
    for (ll j = 0; j <= i - 1; j++) {
      ll now = (ncr[i - 1][j] * exp) % MOD2;
      now = (now * p[n - j]) % MOD2; // 現在の j での値を now として求める
      ans = (ans + now) % MOD2; // ans の値の更新
    }

    cout << ans; // ans の値の出力
    if (i == n) {
      cout << endl; // i = n のときは改行する
    } else {
      cout << " "; // i = n でないときは空白区切りで表示する
    }
  }

  return 0;
}

提出したコード: https://atcoder.jp/contests/abc333/submissions/48837340
GitHub

感想

時間こそかかりましたが、自力で解くことができてとても嬉しいです。

公式解説 *2 ではもっと簡単に解説されていますが、自分の思いついた方法はこれだったので、このように記載しています。

いざコンテストにてこの問題が出たときに、おそらくこの方法で考えていたので、制限時間に間に合うかの勝負になっていたと思います。

*1:https://ja.wikipedia.org/wiki/%E3%83%95%E3%82%A7%E3%83%AB%E3%83%9E%E3%83%BC%E3%81%AE%E5%B0%8F%E5%AE%9A%E7%90%86

*2:https://atcoder.jp/contests/abc333/editorial/7948

yutasの競技プログラミング勉強帖

競技プログラミングの問題についての解説記事を主に書いています。

F - Bomb Game 2 / トヨタ自動車プログラミングコンテスト2023#8 (AtCoder Beginner Contest 333)

問題

補足

入力

条件

解法

ソースコード

感想