ABC471 E - Sum of Square of Sum

考え方

各回スコアの式を展開すると、全て $A_i\times A_j$ という形になる。
全取り出し方の合計は、それぞれの係数がどうなるか考えることで処理すればよい。

まず、$i$ と $j$ が等しい項 $A_i^2$ について。
これは、全取り出し方のうち、$i$ 番が含まれる個数が係数となる。
つまり、$i$ 番以外の $K-1$ 個の選び方 ${}_{N-1}\mathrm{C}_{K-1}$ 通り。
全 $i$ での合計は、全ての $2$ 乗和を愚直計算し、${}_{N-1}\mathrm{C}_{K-1}$ を掛ければよい。

そして、$i$ と $j$ が異なる項 $A_i\times A_j$ について。
$i$ と $j$ を入れ替えたものもまとめて扱うことにして、$i<j$ 前提で考える。
これは、全取り出し方のうち、$i$ 番と $j$ 番が含まれる個数の $2$ 倍が係数となる。
つまり、$i$ 番と $j$ 番以外の $K-2$ 個の選び方 ${}_{N-2}\mathrm{C}_{K-2}$ 通り。
全 $i,j$ での合計は、全ての積の和を計算し、$2{}_{N-2}\mathrm{C}_{K-2}$ を掛ければよい。

しかし、この全ての積の和の方は、愚直に計算すると $O(N^2)$ かかって TLE となる。
これは高校数学でよくあるテクニックが役に立つ。
$1$ 乗和の $2$ 乗から $2$ 乗和を引くと、全ての異なる要素の積の $2$ 倍となる。
$2$ 倍は既にされているので、${}_{N-2}\mathrm{C}_{K-2}$ だけかけることになる。

これによって高速化ができ、全体の計算量は $O(N+\log \mathrm{MOD})$ となる。

入力例1での動作

入力を受け取る。

n: 3
k: 2
a: {1, 10, 100}

まず、$1$ 乗和は $1+10+100=111$、$2$ 乗和は $1^2+10^2+100^2=10101$ となる。

$i$ と $j$ が異なる項の合計については、$111^2-10101=2220$ となる。
また、${}_{N-2}\mathrm{C}_{K-2}={}_{1}\mathrm{C}_{0}=1$ なので、この部分から $2220$ を加える。

$i$ と $j$ が等しい項については、${}_{N-1}\mathrm{C}_{K-1}={}_{2}\mathrm{C}_{1}=2$ なので、$10101\times2=20202$ を加える。

よって、$2220+20202=22422$ となる。

注意点

$A_i^2$ は、int 型からはみ出る。
long long 型を用いること。

答えは $998244353$ で割った余りを要求されているので、剰余類環の考えに従って処理する。
何か足し算や掛け算をするたびに結果を % 998244353 し、割り算は $998244353-2$ 乗したものを掛ける。
二項係数の計算で逆元を求めるため、繰り返し二乗法のアルゴリズムも用意しておくこと。

別解

特になし。