ABC464 F - Random Vault Heist

考え方

まず、$N \leq 20$ の場合の $O(2^N)$ 解法を考える。

例として、金庫が $3$ 個で、中身が $2$ と $3$ と $4$、目標金額が $7$ である場合を考える。
金庫を開ける順序は $6$ 通りあり、以下の寄与を合計して、期待値は $25/3$ となる。

これを、何個目に開ける金庫からの寄与かで和の取り方を作り直す。

さらに、それぞれの中で、直前までにどの金庫を開けたのかで分解する。

それぞれの分子の和は、残っている金庫の総額である。
つまり、全金庫の総額から、既に開けた金庫の総額を引いて求められる。
分母は、「事前に開けたものがそうなった上で、次の $1$ 個で残りのそれぞれの金庫を開ける確率」の分母。

つまり、今回の期待値は、以下で求まる。

$$
\frac{9-0}{{}_3\mathrm{C}_0 \times 3}+\frac{9-2}{{}_3\mathrm{C}_1 \times 2}+\frac{9-3}{{}_3\mathrm{C}_1 \times 2}+\frac{9-4}{{}_3\mathrm{C}_1 \times 2}+\frac{9-5}{{}_3\mathrm{C}_2 \times 1}+\frac{9-6}{{}_3\mathrm{C}_2 \times 1}=\frac{25}{3}
$$

これは、各項ごとに、選んだ金庫の組の個数と合計額のみから $O(1)$ で高速に計算できる。
（逆数に毎回繰り返し二乗法を用いなくて済むように、階乗や各種逆数を事前計算しておく）
つまり、以下のようなデータを高速に構築できれば、$7$ 未満のデータだけ拾うことで $O(2^N)$ で解ける。

このデータの構築は、以下の方法を用いる。

最初に、このような形のデータを用意する。

ここに、金額 $2$ の金庫の影響を処理する。
つまり、すべてのデータについて、合計額に $2$ を加えて、個数が $1$ 大きいところに付け加える。

次に、金額 $3$ の金庫の影響を処理する。
方法は同様だが、$1$ 個のところはマージソートの要領で昇順に保っておく。
（これは、$N \leq 20$ のときには特に意味がないが、$N \leq 40$ に対応するときに役に立つ）

最後に金額 $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$ 未満になるような全組について、「全総額-その値」の総和を出せばよい。
これは、後者のインデックスを進めながら前者のインデックスを戻すツーポインタ法で求められる。
前者側の累積和も事前に取っておけば、両配列の長さの和に比例する時間で求まる。
これを $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$ を前半と後半に分ける。

前半: {2, 3, 5, 7, 11}
後半: {13, 17, 19, 23, 29, 31}

まず、前半について、選ぶ個数ごとに部分集合和を昇順で列挙する。

後半についても同様に列挙する。
以下では、前半から $2$ 個、後半から $2$ 個を既に開けた場合を例として追う。
後半から $2$ 個選んだ部分集合和は次の通りである。

30, 32, 36, 36, 40, 42, 42, 44, 46, 48, 48, 50, 52, 54, 60

前半から $2$ 個選んだ部分集合和と、その累積和は次のようになる。

部分集合和: {5, 7, 8, 9, 10, 12, 13, 14, 16, 18}
累積和:     {0, 5, 12, 20, 29, 39, 51, 64, 78, 94, 112}

全金庫の総額は $160$ である。
後半側の部分集合和を $k$ とする。
前半側の部分集合和が $60-k$ 未満のものだけが、既に盗んだ合計が $X=60$ 未満という条件を満たす。

例えば、いくつかの $k$ について計算すると次のようになる。

後半から $2$ 個選ぶ場合について、全て同様に計算する。
この組み合わせでの分子部分の合計は $10978$ となる。

前半から $2$ 個、後半から $2$ 個の合計 $4$ 個を既に開けている。
その $4$ 個の集合が最初に開けられる確率は $1/{}_{11}\mathrm{C}_4$ である。
次の金庫は残り $7$ 個から一様ランダムに選ばれる。
したがって、この場合の期待値への寄与は次の値である。

$$
10978 \times \frac{1}{{}_{11}\mathrm{C}_4} \times \frac{1}{7}
$$

前半から選ぶ個数と、後半から選ぶ個数の全ての組み合わせについて同じ処理を行う。
それぞれの寄与を剰余類環上で加算すると、最終的な答えは $525950964$ となる。

注意点

A_i、X、部分集合の合計値は、int 型からはみ出る。
long long 型を用いること。

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

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

別解

特になし。