ちょうど \(k+1\) 回目の操作で条件を満たすようになる確率を各 \(1 \leq k \leq N\) について求めます。

\(k\) 回目までに選んだ靴下の色の集合が \(i_1 < i_2 < \dots < i_k\) であり、\(k+1\) 回目に選んだ靴下の色がこのうち \(i_\ell\) であるように選ばれる確率は

\[\displaystyle \frac{k!}{S(S-1)\cdots (S-k)} \left( \prod_{\substack{1 \leq j \leq k \\ j \neq \ell}} A_{i_j} \right) \cdot A_{i_\ell}(A_{i_\ell} - 1)\]

です。これをすべての \(i_1 < i_2 < \dots < i_k\) および \(1 \leq \ell \leq k\) について合計することを考えると、多項式を用いて

\[\frac{k!}{S(S-1)\cdots(S-k)} [x^k y] \prod_{i=1}^N \left(1 + A_ix + A_i(A_i-1)xy\right)\]

と表現できます。

結局、\(2\) 変数多項式の総積 \(\displaystyle \prod_{i=1}^N \left(1 + A_ix + A_i(A_i-1)xy\right)\) を、\(y\) については \(1\) 次までで打ち切ったもの（\(\bmod \ {y^2}\) をとったもの）を求めればよいです。

\(2\) 変数多項式を、\(x\) の多項式 \(f_0, f_1\) のペアを用いた \(f_0 + f_1y\) の形で管理します（pair<vector<mint>, vector<mint>> のようなことです）。

\[(f_0 + f_1y)(g_0 + g_1y) = f_0g_0 + (f_0g_1 + f_1g_0) y \pmod {y^2}\]

という式にもとづき、通常の多項式の総積と同様の方法（分割統治法など）で計算することができます。時間計算量は同じく \(O(N (\log N)^2)\) です（ただし畳み込みの回数が増えており、定数倍は増えます）。

ところで、

\[\frac{f_1}{f_0} + \frac{g_1}{g_0} = \frac{f_0g_1 + f_1g_0}{f_0g_0}\]

であるので、上記のアルゴリズムは有理式の総和を求めるアルゴリズムと一致します。