Sum of Average の解説

Akに対する寄与を考える。

$$\sum_{k=1}^{N} A_k \times \sum_{l=1}^{k} \sum_{r=k}^{N} \frac{1}{r - l + 1}$$

Akを固定した時にここを高速に求めることを考える

$$\sum_{l=1}^{k} \sum_{r=k}^{N} \frac{1}{r - l + 1}$$

ここで、以下のように定義すると

$$H_i = \sum_{k=1}^{i} \frac{1}{k}$$ $$\sum_{l=1}^{k} (H_{N - l + 1} - H_{k - l})$$

とできる。

$$\sum_{l=1}^{k} H_{N - l + 1}$$

についてはj = l - 1とすることで

$$\sum_{j=0}^{k - 1} H_{N - j}$$

とでき、

$$\sum_{l=1}^{k} H_{k - l}$$

については j = k - lとすることで

$$\sum_{j=0}^{k - 1} H_{j}$$

とできるので、

もとの式は

$$\sum_{k=1}^{N} A_k \times \sum_{j=0}^{k - 1} (H_{N - j} - H_{j})$$

とできる。

右側はkをなめると同時に計算することができるのでO(N)で解ける。

fn main() {
    input! {
        n: usize,
        a: [ModInt998244353; n]
    }
    let h = {
        let mut v = vec![];
        v.push(ModInt998244353::new(0));
        for num in 1..=n {
            v.push(ModInt998244353::new(1) / num);
        }
        PrefixSum1D::new(&v)
    };
    let mut ans = ModInt998244353::new(0);
    let mut acc = ModInt998244353::new(0);
    for k in 0..n {
        acc += h.prefix_sum(n - k + 1) - h.prefix_sum(k + 1);
        ans += a[k] * acc;
    }
    println!("{}", ans);
}