Revenge of "The Salary of AtCoder Inc." の解説

dp[i] := x = iの状態からもらえる給料の期待値 としてdpする。

本当にこれをシミュレーションするとさすがに間に合わない。 なので逆から順に考える。

(i + 1..)の期待値とその目が出たときの給料それぞれの区間和が取得できればよいので、fenwicktree or segtreeを用いてx=nの場合から更新していけば求まる。

fn main() {
    input! {
        n: usize,
        a: [usize; n]
    }
    let a = a.iter().map(|&e| ModInt998244353::new(e)).collect_vec();
    let mut seg = Segtree::<M>::new(n + 2);
    for i in 0..n {
        seg.set(i + 1, a[i]);
    }
    let mut dp = Segtree::<M>::new(n + 2);
    for i in (0..=n).rev() {
        let next = seg.prod(i + 1..);
        let pre = dp.prod(i + 1..);
        dp.set(i, (next + pre) / n);
    }
    let ans = dp.get(0);
    println!("{}", ans);
}