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 を取る場合と、23 の両方を取る場合が追加される。

result[0]: {0}
result[1]: {2, 3}
result[2]: {5}

次に値 5 を処理すると、例えば、$1$ 個取る合計に 5 が追加される。
また、$2$ 個取る合計には 2+5=73+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 未満なので、id10 となる。

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$ 未満のものだけを使うので、id9 となる。

後半側の合計 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_iX、部分集合の合計値は、int 型からはみ出る。
long long 型を用いること。

sum1 は部分集合和の総和を持つため、そのまま加算すると long long 型からもはみ出る。
各加算ごとに $998244353$ で割った余りを取ること。

答えは $998244353$ で割った余りを要求されているので、剰余類環の考えに従って処理する。
何か足し算や掛け算をするたびに結果を % 998244353 する。
割り算は通常 $998244353-2$ 乗したものを掛けるが、それでは計算量が増えてしまう。
階乗の事前計算で逆数も求めておき、繰り返し二乗法を毎回しなくて済むように高速化しておくこと。

別解

特になし。