Sum of Square of Sum の解説

$$(a + b + c + d)^2 = a^2 + b ^ 2 + c ^ 2 + d ^ 2 + 2 * ab + 2 * ac + 2 * ad + 2 * bc + 2 * bd + 2 * cd$$

↑前提としてこれが頭の中にある。

まず、ある数Aiが$\binom{n}{k}$の中に含まれる場合の数は、$\binom{n - 1}{k - 1}$ 通り。これだけ $Ai^2$は足し込まれる。

あとは2 * Ai * Ajが何回足し込まれるか。 これはこの二つを除いたk - 2個をn - 2個の中から選べば良いので$\binom{n - 2}{k - 2}$回だけ足されるのでこれらを足せばもとまる。

fn main() {
    input! {
        n: usize,
        k: usize,
        a: [ModInt998244353; n]
    }
    let sum = a.iter().sum::<ModInt998244353>();
    let sum2 = a.iter().map(|&e| e * e).sum::<ModInt998244353>();
    let comb = Comb::<ModInt998244353>::new(n + k + 10);
    if k == 1 {
        return println!("{}", sum2);
    }
    let ans = comb.nck(n - 1, k - 1) * sum2 + (sum * sum - sum2) * comb.nck(n - 2, k - 2);
    println!("{}", ans);
}