ポケポケっぽいAtCoderの問題を数式ゴリゴリで説明をする

# はじめに

[E - Expansion Packs](https://atcoder.jp/contests/abc382/tasks/abc382_e) の問題は、最近の流行りのカードゲームを題材にしているように思えます。

- [トップページ｜ 『Pokémon Trading Card Game Pocket（ポケポケ）』公式サイト](https://www.pokemontcgpocket.com/ja/)


肝心な問題の内容は、確率論とアルゴリズムをうまく組み合わせる必要がある問題です。

自分の出身畑は、数学なので、この問題を数学的に解説してみたいと思います。

(本記事は統計学のアドベントカレンダーの一環として投稿していますが、確率論的な式変形が多いです。一方で統計学でも込み入った式変形をする際は、確率論の知識が必要になることが多いため、この記事は統計学の式変形の参考になるかもしれません。)

---


問題の説明はせずに、
解説から始めます。

問題を読んでいない方は、
以下のリンクから
問題を確認してください。

[E - Expansion Packs](https://atcoder.jp/contests/abc382/tasks/abc382_e)


では、解説を始めます。

---


## 解説

この問題を解くには、以下の2つのステップが必要です。

1. パックを開封したときのレアカードの数の確率分布を求める。
2. 開封するパックの個数の期待値を求める。

一つずつ見ていきます。

### 1. パックを開封したときのレアカードの数の確率分布を求める。

1パックあたり $N$ 枚のカードがあり、それぞれのカード $i$ は確率 $p_i = \frac{P_i}{100}$ でレアカードとなります。  
パックを開けたとき、パック内のレアカード数を表す確率変数 $U$ を考えます。  
ここで $U = \sum_{i=1}^N Y_i$ とし、$Y_i$ は「カード $i$ がレアであるかどうか」を示す指示確率変数（Bernoulli変数）とします。すなわち、  

- $Y_i$: $i$ 番目のカードがレアであるかを表す確率変数。  
  ```math
  Y_i =
  \begin{cases}
    1 & \text{確率 } p_i \text{ でレア}\\
    0 & \text{確率 } 1-p_i \text{ でレアでない}
  \end{cases}
  ```

 ここで $Y_i$ は Bernoulli 分布に従う指示的な確率変数であり、値は 0 または 1 のみを取ります。

- $U_i$: 最初から $i$ 番目までの $i$ 枚のカードを開けたときに手に入るレアカードの総数を表す確率変数。  
  ```math
  U_i = \sum_{j=1}^i Y_j
  ```

このとき、$U_i$ は $Y_1,\dots,Y_i$ の和なので、取りうる値は $0, 1, \dots, i$ となります。

初期条件として、カードを一枚も開けていない状態 ($i=0$) では、レアカードは当然 0 枚なので、  
```math
P(U_0=0)=1,\quad P(U_0=k)=0\,(k \neq 0)
```

#### 分布を段階的に求める

1枚目を開けたとき ($i=1$) は、  
```math
P(U_1=0)=1-p_1,\quad P(U_1=1)=p_1,\quad P(U_1=k)=0\,(k\notin\{0,1\}).
```

2枚目を開けたとき ($i=2$) について考えます。$U_2 = Y_1 + Y_2$ です。  
$Y_2$ は 0 か 1 で、確率はそれぞれ $1-p_2$ と $p_2$ です。

ここで、$U_2$ がちょうど $k$ となるためには2通りのシナリオがあります。

1. 2枚目のカードがレア ($Y_2=1$) の場合:  
   この場合、  
   ```math
   U_2 = U_1 + Y_2 = U_1 + 1.
   ```  
   よって、$U_2=k$ となるには、$U_1 = k-1$ でなければなりません。  

2. 2枚目のカードがレアでない ($Y_2=0$) の場合:  
   この場合、  
   ```math
   U_2 = U_1 + 0 = U_1.
   ```  
   よって、$U_2=k$ となるには、$U_1 = k$ でなければなりません。  

これら2つのシナリオを足し合わせると、  
```math
P(U_2 = k) = p_2 P(U_1 = k-1) + (1-p_2) P(U_1 = k).
```

同様の議論は一般の $i$ についても成立します。  
$i$ 枚目を開けた際に、$U_i = U_{i-1} + Y_i$ であり、$Y_i$ は確率 $p_i$ で1、確率 $1-p_i$ で0です。よって、  
```math
P(U_i = k) = p_i P(U_{i-1}=k-1) + (1-p_i)P(U_{i-1}=k).
```

つまり、$P(U_i)$の$i$に対して、逐次的に更新をすることで$U_N$ i.e. $U$ の分布を求めることができます。この計算量は $O(N^2)$ です。

### 2. 開封するパックの個数の期待値を求める。

次に、問題文にある「開封するパックの個数の期待値」を求めるために、以下の確率変数を導入します。

- $Z_i$: 「レアカードがあと $i$ 枚必要な状態」から出発し、目標（あと $i$ 枚以上）を達成するまでに必要な追加パック数を表す確率変数とします。
 ここで「あと $i$ 枚必要」とは、「今手元にあるレアカード数を基準に、さらに $i$ 枚集めるまでやる」という意味です。

明らかに $Z_0 = 0$ で、よって $E[Z_0] = 0$ となります。

#### 期待値の定義

期待値は定義上、確率変数の取りうる全ての値についての重み付き平均です。

```math
E[Z_i] = \sum_{z=0}^{\infty} z \cdot P(Z_i = z).
```

ここで、$Z_i$ は「あと $i$ 枚必要」な状態から始めて、ゴールするまでに何パック開封するか、という分布を持つ確率変数です。つまり、求めたい値は $E[Z_X]$ です。

#### 条件付き期待値を用いた導出

まず、期待値計算でよく使われる手法として、ある補助的な確率変数について条件付き期待値を求め、その後に条件づけを外すことを適用します。最初に開けたパックを基準にして、それ以降のパックを考えていきます。

1. 最初に1パックを必ず開封します。この1パックで得られるレアカードの枚数は確率変数 $U$ に従います。
2. $U$ の値によって、残り必要な枚数が変動します。

ここで、$Z_i$ を期待値で表します。条件付き期待値の定義より、

```math
E[Z_i] = E[E[Z_i \mid U]].
```

$U$ の値ごとに $Z_i$ の期待値を考えると：

- $U=k$ の場合:

  - もし $k \ge i$ なら $Z_i = 1$ よって、$E[Z_i \mid U=k] = 1$。
  - もし $k < i$ なら $Z_i = 1 + Z_{i-k}$、よって、$E[Z_i \mid U=k] = 1 + E[Z_{i-k}]$。

これをまとめると、

```math
E[Z_i \mid U=k] =
\begin{cases}
1 & k \ge i \\
1 + E[Z_{i-k}] & k < i
\end{cases}
```

続いて、期待値の定義から下記の式が成り立ちます。

```math
E[Z_i] = \sum_{k=0}^{N} P(U=k) E[Z_i \mid U=k].
```

この右辺で、$k \ge i$ の項では $E[Z_i \mid U=k]=1$、  
$0 \le k < i$ の項では $E[Z_i \mid U=k]=1+E[Z_{i-k}]$ となります。したがって、

```math
E[Z_i] = \sum_{k=0}^{i-1} P(U=k) (1 + E[Z_{i-k}]) + \sum_{k=i}^{N} P(U=k) \cdot 1.
```

ここで、$\sum_{k=0}^{N} P(U=k) = 1$ なので、整理すると、

```math
E[Z_i] = 1 + \sum_{k=0}^{i-1} P(U=k) E[Z_{i-k}].
```

上記の式を変形すると、

```math
E[Z_i] = \frac{1}{1-P(U=i)} \left(1 + \sum_{k=1}^{i-1} P(U=k) E[Z_{i-k}]\right).
```

ここで $E[Z_0]=0$ の初期条件を用いて、$i=1$ から順に計算することで $E[Z_X]$ を求めることができます。$U\leq N$であるため、計算量は$O(X \min(X, N))$ となります。

---

## 感想

１つ目の分布の計算については、発想さえあれば漸化式の容量で解ける問題でした。

一方、２つ目の計算については、条件付き期待値を使った期待値の計算が必要で、確率変数の設計が必要で、少し難易度が高かったです。ただ、一つのパックに注目して考えると、意外と自然な発想なようにも思えます。

この手の問題を時間内に解くため、日々確率だとか期待値の計算に触れておくことが大事だと感じました。