ARC034 D - インフレゲーム

問題

atcoder.jp

解法

期待値ではなく,スコアの総和を考えることにします.スコアの総和 $S$ が求まれば,答えは $\dfrac{ S } { ( A + B + C ) !}$ として容易に計算されます.

まずは具体的な例を考えます.

$a _ 1, a _ 2, b _ 1, b _ 2, a _ 3, b _ 3, c _ 1$ の順にカードを引いたとき,スコアは $( ( a _ 1 + a _ 2 ) b _ 1 b _ 2 + a _ 3 ) b _ 3 = a _ 1 \cdot b _ 1 b _ 2 b _ 3 + a _ 2 \cdot b _ 1 b _ 2 b _ 3 + a _ 3 \cdot b _ 3$ となります.従って,$a _ i$ の係数が $b _ { x _ 1 } \cdots b _ { x _ k } \; (1 \leq x _ 1 \lt \cdots \lt x _ k \leq B)$ となるようなカードの並べ方の数を $f _ i (x _ 1, \ldots, x _ k)$ とおけば,スコアの総和 $S$ は次で求められます.

$$ S = \sum _ { i = 1 } ^ A \sum _ { 1\leq x _ 1 \lt \cdots \lt x _ k \leq B } a _ i \cdot f _ i (x _ 1, \ldots, x _ k) \cdot \prod _ {j = 1} ^ k b _ { x _ j } $$

$f _ i (x _ 1, \ldots, x _ k)$ の具体的な表示について考えます.

まず,山札において,$a _ i$ の後に $b _ { x _ 1 }, \ldots, b _ { x _ k }$ が順不同で並んでおり,その後に $c _ 1, \ldots, c _ C$ が順不同で並んでいる必要があります.

ここで,$a _ i$ の並べ方は $1$ 通り,$b _ { x _ 1 }, \ldots, b _ { x _ k }$ の並べ方は $k !$ 通り,$c _ 1, \ldots, c _ C$ の並べ方は $C !$ 通りです.

$b _ { x _ 1 }, \ldots, b _ { x _ k }$ や $c _ 1, \ldots, c _ C$ はどのように並んでいても等価なので,列 $(a _ i, b _ { x _ 1 }, \ldots, b _ { x _ k }, c _ 1, \ldots, c _ C)$ に対して残りの $a$ や $b$ のカードを挿入する方法の数を考えます.

まだ列に含まれていない $b$ のカードは $B - k$ 枚あり,これらは $a _ i$ よりも左か,または $c _ 1$ よりも右に並べられている必要があります.従って,$a _ i, b _ {x _ 1}, \ldots, b _ {x _ k}, c _ 1$ をまとめて $1$ 枚のカードとして見ることで,まだ列に含まれていない $B - k$ 枚の $b$ のカードを挿入する方法は $\displaystyle \binom{B - k + C}{C} \cdot (B - k) ! = \dfrac{(B - k + C) !}{C !}$ 通りあることが分かります.

残りの $A - 1$ 枚の $a$ のカードは $a _ i$ の係数には影響しないのでどこに挿入してもよく,即ち $\displaystyle \binom{A + B + C}{A - 1} \cdot (A - 1) ! = \dfrac{(A + B + C) !}{(B + C + 1) !}$ 通りの挿入のしかたがあります.

以上より,次の $f _ i (x _ 1, \ldots, x _ k)$ の表示を得ます.重要な事実として,$i$ に依存しないことを注意しておきます.

$$ \begin{aligned} f _ i (x _ 1, \ldots, x _ k) &= k! \cdot C! \cdot \dfrac{(B - k + C) !}{C !} \cdot \dfrac{(A + B + C) !}{(B + C + 1) !} \\ &= \dfrac{(A + B + C) !}{B + C + 1} \cdot \binom{B + C}{k} ^ {-1} \end{aligned} $$

従って,答えは次のように計算されます.

$$ \begin{aligned} \dfrac{S}{(A+B+C)!}&=\dfrac{1}{B + C + 1} \sum _ { i = 1 } ^ A \sum _ { 1\leq x _ 1 \lt \cdots \lt x _ k \leq B } a _ i \cdot \binom{B + C}{k} ^ {-1} \cdot \prod _ {j = 1} ^ k b _ { x _ j } \\ &= \dfrac{1}{B + C + 1} \sum _ {i = 1} ^ A a _ i \sum _ {k = 0} ^ B \binom{B + C}{k} ^ {-1} \cdot T _ k \end{aligned} $$

途中で $\displaystyle T _ k = \sum _ { 1\leq x _ 1 \lt \cdots \lt x _ k \leq B }\prod _ {j = 1} ^ k b _ { x _ j }$ としました.$T _ k$ は,$\displaystyle F(x) := \prod _ {i = 1} ^ B (1 + b _ i x)$ の $x ^ k$ の係数として計算されます.$F(x)$ はナイーブな DP で $\Theta(B ^ 2)$ で,あるいは FFT を用いることで $\Theta(B (\log B) ^ 2)$ で計算できます.

$\displaystyle \sum _ {i = 1} ^ A a _ i$ の計算は $\Theta(A)$ で,式中に現れる二項係数は $\Theta(B)$ で前計算可能なので,本問を $\Theta(A + B ^ 2)$ や $\Theta(A + B (\log B) ^ 2)$ で解くことが出来ました.

感想

現代なら $\mathrm{mod}\; 998244353$ 出力で $1 \leq A, B, C\leq 2\cdot 10 ^ 5, 0 \leq a _ i, b _ i \lt 998244353$ みたいな制約になりそう.$C$ は $B + C + 1 \lt 998244353$ の範囲で大きくできるけど,制約が不自然になるから難しそう (やるなら $A + B + C \lt 998244353$ とか?).