目次
第七回日本最強プログラマー学生選手権~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$
- 入力される値はすべて整数
解法
Fにしてはなかなか考察+実装の難度が高め。
大まかな方針
最後に空ける金庫 $i$ を固定する。
それ以外から $k$ 個を選んで合計金額が $[X-A_i,X)$ 円の範囲になるような(順序を気にしない)選び方1つに対し、 実際に強盗がそのような金庫を開け、かつ最後に $i$ を開けて終わることになる確率は、$k$ だけに依存する。
- $P(k)=\dfrac{k!(N-1-k)!}{N!}$
よって、そのような選び方の個数を $t_{i,k}$、金額の総和を $m_{i,k}$ とすると、$(i,k)$ の組に対する期待値の寄与は
- $E(i,k)=(A_i \cdot t_{i,k} + m_{i,k})P(k)$
となる。全ての $(i,k)$ でこれを合計したものが答えとなる。
高速化
制約が半分全列挙と言っている。
金庫を半分に分け、集合 $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})$ で答えが求められる。

