分割統治及びNTTを用いることは公式解説と同様です。

$K$ 回目で初めて二枚目が揃う靴下の取り出し方が何通りかを求めることが可能です。

余事象を用いない解法

靴下 $i$ を二足目として引くことを固定した場合に $K$ 回目で二枚目を引くことができる選び方は、 $i$ の一枚目と残りの $K-2$ 枚の靴下のどれを選ぶか否かを選んだ後に $(K-1)!$ で並べ替え、最後に二枚目である $A_i-1$ 通りを選ぶ事を考えることで以下の通りに表現できます。

\[(K-1)![x^K]\left( A_i(A_i-1)x^2\prod_{j=1}^{i-1}(1+A_jx)\prod_{j=i+1}^{N}(1+A_jx) \right)\]

従って $f(L,R)=\sum_{i=L}^{R} A_i(A_i-1)x^2\prod_{j=L}^{i-1}(1+A_jx)\prod_{j=i+1}^{R}(1+A_jx)$ とした際の $f(1,N)$ を求めることを目標とします。

各種 $\prod(1+A_jx)$ を一括化したいことを考えると、次の式が成立します。

\[f(L,R) = \begin{cases} A_L(A_L-1)x^2 & (L=R)\\ f(L,M)\prod_{j=M+1}^R(1+A_jx)+ f(M+1,R)\prod_{j=L}^M(1+A_jx) & (L<R) \end{cases} \]

これらは $f(L,R)$ 及び $\prod_{j=L}^R(1+A_jx)$ を分割統治で同時に算出していくことで $f(1,N)$ を $O(N\log^2N)$ で求めることが可能になります。

残るは順次二枚目の靴下が揃っていない場合の余事象を計算してゆくことで $O(N)$ で答えることが可能です。

余事象を用いる解法

数え上げをすること自体は公式解説で出てくる係数の余事象を用いると容易に可能です。

公式解説同様、 $p(x)=\prod_{j=1}^N(1+A_jx)$ の導出をすると各 $x^K$ の係数は「 $K$ 枚引いた時に、まだ靴下が揃っていない時の靴下の引き方」になり、 $K!$ を乗算することで並べ方の列挙になります。

ここで「 $K-1$ 枚を引いて揃っていない並べ方 $(K-1)![x^{K-1}]p(x)$ 」に対して「 $(\sum A_i)-K+1$ 枚のいずれか引く」という行為を乗算として扱い 「 $K$ 枚を引いて揃っていない並べ方 $K![x^K]p(x)$ 」を引くことでも同様の値が出ます。

詳細は提出回答の48行目のassertをご覧ください。

https://atcoder.jp/contests/abc352/submissions/53162179