目次
第七回日本最強プログラマー学生選手権~Advance~ -予選- (AtCoder Beginner Contest 464)F問題メモ
F - Random Vault Heist
問題文
- コワスギ銀行には $N$ 個の金庫があります。金庫 $i$ には $A_i$ 円が入っています。
- ある日、強盗がコワスギ銀行に入りました。強盗は盗んだ金額の合計が $X$ 円以上になるまで、以下の操作を繰り返します。
- まだ開けていない金庫の中から一様ランダムに $1$ つ選んで開けて、中のお金をすべて盗む。
- 強盗によって盗まれた金額の合計の期待値を $\text{mod }998244353$ で求めてください。
制約
- $1 \le N \le 40$
- $1 \le A_i \le 10^{16}$
- $1 \le X \le \sum_{i=1}^{N} A_i$
- 入力される値はすべて整数
解法
単純な設定ながら、なかなか考察+実装の難度が高め。
まず、このように「残っているものからその都度、一様ランダムにアイテムを選んでいく場合の確率・期待値」は、 「順列をランダムに1つ選び、先頭から選んでいく」とした場合の確率・期待値と一致する。 最初にサイコロを振って選ぶ金庫を決めておく、と考えると直感的かもしれない。
大まかな方針
最後に空ける金庫 $i$ を固定する。
$i$ の前に開ける金庫の集合を $S$ とする。
$S$ 内の金庫の合計金額が $[X-A_i,X)$ 円の範囲であれば、実際に最後に $i$ を開けて終わることになる。
ある $S$ に対し、$S$ に含まれる金庫を全て開け、次に $i$ を開けるような確率は、要素数 $k=|S|$ だけに依存する。 先頭 $k$ 個に $S$ 内の金庫を順序を問わず配置し、$k+1$ 番目が $i$ で、残り $N-1-k$ 個はまた順序を問わず配置されるような順列の個数がわかればよい。
- $P(k)=\dfrac{k!(N-1-k)!}{N!}$
$(i,k)$ に対し、合計金額が $[X-A_i,X)$ 円の範囲となる $k$ 個の選び方の個数を $t_{i,k}$、その金額の総和(※)を $m_{i,k}$ とすると、$(i,k)$ の組に対する期待値の寄与は
- $E(i,k)=(A_i \cdot t_{i,k} + m_{i,k})P(k)$
となる。全ての $(i,k)$ でこれを合計したものが答えとなる。
(※)金額の総和とは、「$k$ 個で合計金額が $[X-A_i,X)$ の範囲内」という選び方の“合計金額”を、条件を満たす選び方全てで合計したもの。$A_i$ は入ってない。
高速化
制約が半分全列挙と言っている。
金庫を半分に分け、集合 $L,R$ とする。最後に空ける金庫 $i$ は $L$ の中から1つ選ぶとする。($L,R$ を入れ替えて2通りやる)
$R$ からの選び方を $2^{|R|}$ 通り全列挙する。個数 $k_R$ ごとに分類し、昇順にソートしておく。
$B[k_R]$ を「$R$ から $k_R$ 個選んだ時に取りうる金額のリスト(昇順ソート済み)」とする。
Rに含まれる金庫の金額 = { 1, 10, 100 }
kR 取り得る金額のリスト
0 B[0] = [ 0 ]
1 B[1] = [ 1, 10, 100 ]
2 B[2] = [ 11, 101, 110 ]
3 B[3] = [ 111 ]
$i$ を除く $L$ からも同様に $2^{|L|-1}$ 通り全列挙する。
こちらは、個数情報は残しつつも(個数毎では無く)全体を降順にソートしておく。これを $C$ とする。
i以外のLに含まれる金庫の金額 = { 2, 20, 200 }
(金額, 個数) の降順リスト
C = [ (222,3), (220,2), (202,2), (200,1), (22,2), (20,1), (2,1), (0,0) ]
$k_R$ を固定する。その後、$C$ を先頭から見ていく。
$L$ から (金額, 個数)$=(m_L,k_L)$ 個選んだ時に、 $m_L+m_R$ が $[X-A_i,X)$ の範囲に収まるような $R$ 側の金額 $m_R$ は、 $B[k_R]$ 上で連続した区間に並ぶ。
そのような $B[k_R]$ 上のindex区間を $[a,b)$ とすると、
- $E(i,k_L+k_R)=((A_i + m_L) \cdot (b-a) + \sum_{j=a}^{b-1}B[k_R][j])P(k_L+k_R)$
として、期待値が求められる。(区間和の部分は累積和を事前計算しておく)
$C$ を降順に、$B[k_R]$ を昇順に並べたため、このような区間 $[a,b)$ は、
$C$ を進めて行くにあたり尺取り法が適用でき、
1つの $k_R$ に対して $O(|C|\log{|C|} + |B[k_R]|\log{|B[k_R]|})$ で、$k_R$ に対する答えが全て求められる。
$k_R$ を全通り試すと $O(2^{|L|}|L|+2^{|R|}|R|) = O(N2^{N/2})$ となる。
更に外側で $i$ を全通り試すので、全体は $O(N^22^{N/2})$ で答えが求められる。

