ABC464 F - Random Vault Heist
ランダム強盗
考え方
まず、$N \leq 20$ の場合の $O(2^N)$ 解法を考える。
例として、金庫が $3$ 個で、中身が $2$ と $3$ と $4$、目標金額が $7$ である場合を考える。
金庫を開ける順序は $6$ 通りあり、以下の寄与を合計して、期待値は $25/3$ となる。
| 開ける順序 | 期待値の寄与 |
|---|---|
| ${2,3,4}$ | $(2+3+4)/6 = 3/2$ |
| ${2,4,3}$ | $(2+4+3)/6 = 3/2$ |
| ${3,2,4}$ | $(3+2+4)/6 = 3/2$ |
| ${3,4}$ | $(3+4)/6 = 7/6$ |
| ${4,2,3}$ | $(4+2+3)/6 = 3/2$ |
| ${4,3}$ | $(4+3)/6 = 7/6$ |
これを、何個目に開ける金庫からの寄与かで和の取り方を作り直す。
| 何個目の金庫か | 期待値の寄与 |
|---|---|
| $1$ 個目の金庫 | $(2+2+3+3+4+4)/6 = 3$ |
| $2$ 個目の金庫 | $(3+4+2+4+2+3)/6 = 3$ |
| $3$ 個目の金庫 | $(4+3+4+3)/6 = 7/3$ |
さらに、それぞれの中で、直前までにどの金庫を開けたのかで分解する。
| 何個目の金庫か | 直前までに開けた金庫 | 期待値の寄与 |
|---|---|---|
| $1$ 個目の金庫 | $( )$ | $(2+3+4)/3 = 3$ |
| $2$ 個目の金庫 | $(2)$ | $(3+4)/6 = 7/6$ |
| $2$ 個目の金庫 | $(3)$ | $(2+4)/6 = 1$ |
| $2$ 個目の金庫 | $(4)$ | $(2+3)/6 = 5/6$ |
| $3$ 個目の金庫 | $(2,3)$ | $4/3 = 4/3$ |
| $3$ 個目の金庫 | $(2,4)$ | $3/3 = 1$ |
それぞれの分子の和は、残っている金庫の総額である。
つまり、全金庫の総額から、既に開けた金庫の総額を引いて求められる。
分母は、「事前に開けたものがそうなった上で、次の $1$ 個で残りのそれぞれの金庫を開ける確率」の分母。
つまり、今回の期待値は、以下で求まる。
$\dfrac{9-0}{{}_3\mathrm{C}_0 \times 3}+\dfrac{9-2}{{}_3\mathrm{C}_1 \times 2}+\dfrac{9-3}{{}_3\mathrm{C}_1 \times 2}+\dfrac{9-4}{{}_3\mathrm{C}_1 \times 2}+\dfrac{9-5}{{}_3\mathrm{C}_2 \times 1}+\dfrac{9-6}{{}_3\mathrm{C}_2 \times 1} = \dfrac{25}{3}$
これは、各項ごとに、選んだ金庫の組の個数と合計額のみから $O(1)$ で高速に計算できる。
(逆数に毎回繰り返し二乗法を用いなくて済むように、階乗や各種逆数を事前計算しておく)
つまり、以下のようなデータを高速に構築できれば、$7$ 未満のデータだけ拾うことで $O(2^N)$ で解ける。
| 選んだ金庫の個数 | 合計額 |
|---|---|
| $0$ 個 | $0$ |
| $1$ 個 | $2, 3, 4$ |
| $2$ 個 | $5, 6, 7$ |
| $3$ 個 | $9$ |
このデータの構築は、以下の方法を用いる。
最初に、このような形のデータを用意する。
| 選んだ金庫の個数 | 合計額 |
|---|---|
| $0$ 個 | $0$ |
| $1$ 個 | |
| $2$ 個 | |
| $3$ 個 |
ここに、金額 $2$ の金庫の影響を処理する。
つまり、すべてのデータについて、合計額に $2$ を加えて、個数が $1$ 大きいところに付け加える。
| 選んだ金庫の個数 | 合計額 |
|---|---|
| $0$ 個 | $0$ |
| $1$ 個 | $2$ |
| $2$ 個 | |
| $3$ 個 |
次に、金額 $3$ の金庫の影響を処理する。
方法は同様だが、$1$ 個のところはマージソートの要領で昇順に保っておく。
(これは、$N \leq 20$ のときには特に意味がないが、$N \leq 40$ に対応するときに役に立つ)
| 選んだ金庫の個数 | 合計額 |
|---|---|
| $0$ 個 | $0$ |
| $1$ 個 | $2, 3$ |
| $2$ 個 | $5$ |
| $3$ 個 |
最後に金額 $4$ の金庫の影響を処理すれば、前述のデータが $O(2^N)$ で構築できる。
あとは、$1$ つずつ分数計算をすればよい。
以上が、$N \leq 20$ の場合に $O(2^N)$ で解く方法である。
さて、ここまでは前座で、この問題はここからが本番である。
$N \leq 40$ の場合には $O(2^N)$ では間に合わず、$O(N 2^{N/2})$ くらいでどうにかしなければならない。
まず、$2^40$ 通りの列挙は絶対に無理なので、半分全列挙で前後半それぞれの $2^20$ 通りずつ全列挙する。
そして、前半から $i$ 個、後半から $j$ 個とるパターンごとにまとめて処理していく。
これらの寄与を計算するときの分母は、${}_N\mathrm{C}_{i+j} \times (N-i-j)$ でよい。
分子側は前半から $i$ 個とるデータの配列と、後半から $j$ 個とるデータの配列を見て求める。
前者の値と後者の値の和が $X$ 未満になるような全組について、「全総額-その値」の総和を出せばよい。
これは、後者のインデックスを進めながら前者のインデックスを戻すツーポインタ法で求められる。
前者側の累積和も事前に取っておけば、$O(両配列の長さの和)$ で求まる。
これを $i$ と $j$ 全ての組について実行して全ての和をとればよい。
半分全列挙のために A を分けたときに、左側に $l$ 個、右側に $r$ 個あったとする。
すると、ツーポインタ法の二重ループ全体で $O(l2^r+r2^l)$ 回かかる。
これは $l$ と $r$ を均等に分けたら $O(N 2^{N/2})$。
他にこれより計算量が膨らむところはないので、全体の計算量も $O(N 2^{N/2})$。
これで $N \leq 40$ の場合にも解ける。
入力例3での動作
入力例 $1$ があまり動作説明に適さないので、入力例 $3$ で例を示す。
入力を受け取る。
n: 11
x: 60
a: {2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31}
配列 a を前半と後半に分ける。
b: {2, 3, 5, 7, 11}
c: {13, 17, 19, 23, 29, 31}
前半から $k$ 個取ったときの合計を、$k$ ごとに列挙する。
最初は、$0$ 個取った合計だけがある。
result[0]: {0}
result[1]: {}
result[2]: {}
result[3]: {}
result[4]: {}
result[5]: {}
値 2 を処理すると、2 を取る場合が追加される。
result[0]: {0}
result[1]: {2}
次に値 3 を処理すると、3 を取る場合と、2 と 3 の両方を取る場合が追加される。
result[0]: {0}
result[1]: {2, 3}
result[2]: {5}
次に値 5 を処理すると、例えば、$1$ 個取る合計に 5 が追加される。
また、$2$ 個取る合計には 2+5=7 と 3+5=8 が追加される。
各配列が昇順になるようにマージソートしながら追加する。
result[0]: {0}
result[1]: {2, 3, 5}
result[2]: {5, 7, 8}
result[3]: {10}
以下同様に繰り返すと、前半側の列挙結果は次のようになる。
vv1[0]: {0}
vv1[1]: {2, 3, 5, 7, 11}
vv1[2]: {5, 7, 8, 9, 10, 12, 13, 14, 16, 18}
vv1[3]: {10, 12, 14, 15, 16, 18, 19, 20, 21, 23}
vv1[4]: {17, 21, 23, 25, 26}
vv1[5]: {28}
後半側も同様に列挙する。
今回の二重ループの説明では、後半から $2$ 個取る場合だけを見る。
vv2[2]: {30, 32, 36, 36, 40, 42, 42, 44, 46, 48, 48, 50, 52, 54, 60}
また、a 全体の合計値を取得する。
sum_all: 160
ここから、前半から $2$ 個、後半から $2$ 個取る場合の寄与の計算だけ例示する。
vv1[2]: {5, 7, 8, 9, 10, 12, 13, 14, 16, 18}
vv2[2]: {30, 32, 36, 36, 40, 42, 42, 44, 46, 48, 48, 50, 52, 54, 60}
まず、vv1[2] の累積和 sum1[2] を作る。
sum1[2][t] は、vv1[2] の先頭 $t$ 個の和である。
vv1[2]: {5, 7, 8, 9, 10, 12, 13, 14, 16, 18}
sum1[2]: {0, 5, 12, 20, 29, 39, 51, 64, 78, 94, 112}
vv2[2] の値を $1$ つずつ見る。
最初に、後半から取った合計が 30 の場合を考える。
このとき、前半から取った合計が $60-30=30$ 未満なら、すでに盗んだ集合 $S$ の合計が $X$ 未満である。
vv1[2] の値はすべて 30 未満なので、id は 10 となる。
k: 30
x-k: 30
id: 10
条件を満たす前半側の合計は、次の $10$ 個である。
5, 7, 8, 9, 10, 12, 13, 14, 16, 18
そのため、条件を満たす前半側の合計の総和は、sum1[2][10] を使って $112$ と求められる。
sum1[2][id]: 112
この k=30 について、次に盗む金額の期待値の分子部分をまとめると、次の値になる。
(sum_all-k)*id-sum1[2][id]
= (160-30)*10-112
= 1188
同様に、vv2[2] の各値について処理すると次のようになる。
例えば k=42 の場合は、前半側の合計が $18$ 未満のものだけを使うので、id は 9 となる。
後半側の合計 k |
x-k |
id |
sum1[2][id] |
(sum_all-k)*id-sum1[2][id] |
|---|---|---|---|---|
| 30 | 30 | 10 | 112 | 1188 |
| 32 | 28 | 10 | 112 | 1168 |
| 36 | 24 | 10 | 112 | 1128 |
| 36 | 24 | 10 | 112 | 1128 |
| 40 | 20 | 10 | 112 | 1088 |
| 42 | 18 | 9 | 94 | 968 |
| 42 | 18 | 9 | 94 | 968 |
| 44 | 16 | 8 | 78 | 850 |
| 46 | 14 | 7 | 64 | 734 |
| 48 | 12 | 5 | 39 | 521 |
| 48 | 12 | 5 | 39 | 521 |
| 50 | 10 | 4 | 29 | 411 |
| 52 | 8 | 2 | 12 | 204 |
| 54 | 6 | 1 | 5 | 101 |
| 60 | 0 | 0 | 0 | 0 |
これらの合計は $10978$ である。
1188+1168+1128+1128+1088+968+968+850+734+521+521+411+204+101+0 = 10978
今回は前半から $2$ 個、後半から $2$ 個取っているので、すでに盗んだ金庫の個数は $4$ 個である。
「その $4$ 個が先頭に来る確率」は、二項係数の逆数を用いて処理する。
comb_inv(11,4): 674571184 (実際の分数の値は 1/330)
また、次に盗む金庫は残り $11-4=7$ 個から一様ランダムに選ばれる。
よって、さらに $1/7=855638017$ をかける。
したがって、この場合に結果に加算する値は、剰余類環上では $10978 * 674571184 * 855638017$ である。
同様の処理を、前半から取る個数 i と後半から取る個数 j のすべての組み合わせについて行う。
最後に、結果が負になっていれば $998244353$ を足してから出力する。
注意点
A_i、X、部分集合の合計値は、int 型からはみ出る。
long long 型を用いること。
sum1 は部分集合和の総和を持つため、そのまま加算すると long long 型からもはみ出る。
各加算ごとに $998244353$ で割った余りを取ること。
答えは $998244353$ で割った余りを要求されているので、剰余類環の考えに従って処理する。
何か足し算や掛け算をするたびに結果を % 998244353 する。
割り算は通常 $998244353-2$ 乗したものを掛けるが、それでは計算量が増えてしまう。
階乗の事前計算で逆数も求めておき、繰り返し二乗法を毎回しなくて済むように高速化しておくこと。
別解
特になし。