以降では、\(p_i, q_i\) は入力での値を \(0.01\)倍したものを表すこととします。

概要

結論を先に述べると、次のようなアルゴリズムで期待値が計算できます。

1. \(S = \{ (p_i q_i, i) \}_{1 \leq i \leq N}\) とする。また、\(X = 0, Y = 1\) と初期化する。
2. 以下を \(K\) 回繰り返す。
   1. \(X \leftarrow X + Y\) と更新する。
   2. \(S\) の要素の内、第一要素が最大であるものの一つを \((z, i)\) とする。\(Y \leftarrow Y - z, \, S \leftarrow S \cup \{ (z (1 - q_i), i) \} \setminus \{ (z, i) \}\) と更新する。
3. \(K - X\) を出力する。

優先度付きキューなどを用いて実装すると計算量は \(O((N + K) \log N)\) となり、十分高速です。

上記の手法は、扉が開く確率が最も高い鍵を貪欲に試す方法となっています。 以降では、この方法が期待値的に最適であることを説明します。

期待値の計算

「\(i \, (1 \leq i \leq K)\) 回目より前に扉が開いていないならば、 \(i\) 回目は鍵 \(A_i\) を選ぶ」という方法を戦略 \(A = (A_1, \dots, A_K)\) と呼ぶこととします。

戦略 \(A\) に従って鍵を選んだ場合の残った所持金の期待値は、\(i\) 回目の試行で初めて扉が開く確率を \(d_i\) として、

\[ \sum_{i = 1}^{K - 1} (K - i) \times d_i \tag{1}\]

のように計算できます。ここで、「\(i\) 回目の操作の後に支払いを行う必要があるか？」と考えると、\(s_0 = 1, \,s_i = 1 - d_1 - \cdots - d_i\) として、

\[ K - \left(s_0 + s_1 + s_2 + \cdots + s_{K - 1} \right) \tag{2}\]

のように期待値が計算出来て、式 \((1)\) と一致していることが確認できます。

よって、残った所持金を最大化したい場合、\(s_0 + s_1 + s_2 + \cdots + s_{K - 1}\) を最小化すればよいことが分かります。\(s_i\) は \(i\) 回目までの試行で扉が開いていない確率を表すことから、\(n_{i, j} = \# \{ k \leq i \,|\, A_k = j\}\) として、

\[ s_i = \sum_{j = 1}^N p_j \times (1 - q_j)^{n_{i, j}} \tag{3}\]

と計算できます。

最適な戦略

前述した通り、目的は \(s_0 + s_1 + s_2 + \cdots + s_{K - 1}\) の最小化ですが、ここでは各 \(s_i\) の最小値について考察します。任意の戦略 \(A\) に対して \((s_0, s_1, \dots, s_K)\) は単調減少列であり、減少幅 \(- (s_i - s_{i - 1}) = d_i\) は、

\[d_i = p_{A_i} q_{A_i} \times (1 - q_{A_i})^{n_{i, A_i} - 1} \tag{4}\]

と計算できます。よって、各 \(d_i\) の取りうる値の集合は \(D \coloneqq \{ p_j q_j (1 - q_j)^m \}_{1 \leq j \leq N, 0 \leq m}\) となるため、\(\{d_1, \dots, d_i\}\) が集合 \(D\) の上位 \(i\) 個からなる場合に \(s_i\) は最小値を取ります。

ここで、\((1 - q_j)^m\) が \(m\) に関して単調減少することから、 \(D\) の上位 \(i\) 個は、\(\sum_{j = 1}^N m_j = i\) を満たすある非負整数列 \((m_1, \dots, m_N)\) を用いて、

\[ \{p_1 q_1, \dots, p_1 q_1 (1 - q_1)^{m_1 - 1}, \dots, p_N q_N (1 - q_N)^{m_N - 1} \} \tag{5}\]

と表されます。ゆえに、戦略 \(A\) の \(n_{i, j} = \# \{ k \leq i \,|\, A_k = j\}\) が \(m_j\) に一致する時、\(s_i\) は最小値を取ることが分かります。

以上のような方法で \(s_i\) が最小となる戦略 \(A\) について、\(A_{i + 1}\) を適切な要素とすることで、\(s_{i + 1}\) も最小値を取るようにできます。具体的には、\(D\) で \(i + 1\) 番目に大きい要素を \(p_j q_j (1 - q_j)^m\) として、\(A_{i + 1} = j\) とすればよいです。全ての \(i\) についてこのように \(A_i\) を定めた戦略 \(A\) は、明らかに \(s_0 + s_1 + s_2 + \cdots + s_{K - 1}\) を最小化するので、期待値的に最適な戦略となります。

アルゴリズム

後は、上記の戦略によって選ぶ鍵を決定し、式 \((2)\) に基づいて期待値を計算すればよいです。\(D\) から大きな要素を取り出すには、各 \(j \, (1 \leq j \leq N)\) について \(p_j q_j (1 - q_j)^{n_{i, j}}\) を管理し、その最大値を選択します。最初に述べたアルゴリズムにおいて、\(p_j q_j (1 - q_j)^{n_{i, j}}\) を管理したものが \(S\) であり、各時点での \(s_i\) が \(Y\) となっています。

以上により、この問題を解くことが出来ました。

補足 1

今回は厳密な証明のために \(s_i\) を中心とした解説を行いましたが、今回取った最適戦略は、直感的には「扉が開く条件付き確率が最も大きい鍵を選択する」と解釈できます。実際に、鍵 \(j\) を今まで \(n_{i, j}\) 回試して扉が開かないという条件の下、鍵 \(j\) が本物である確率は、

\[\frac{(j が本物かつ扉が開かない確率)}{(扉が開かない確率)} \propto p_j (1 - q_j)^{n_{i, j}}\]

となるため、鍵 \(j\) を選択して扉が開く確率は \(p_j q_j (1 - q_j)^{n_{i, j}}\) に比例します。

補足 2

この問題と同様の設定で、\(K\) 回以内に扉が開く確率の最大値を求める場合は、\(D\) の上位 \(K\) 個目の値を決め打つ二分探索を行うことで \(O(N \log K)\) の計算量で解くことが出来ます。

原案 : Alt3

解答・改題 : carrot46, cuthbert, noimi