AtCoder Beginner Contest 452 E

https://atcoder.jp/contests/abc452/tasks/abc452_e

jの剰余を使うので、jを固定すると考えます。例えばj=3とすると、

$A_1 + 2A_2 + A_4 + 2A_5 + \cdots$

となるので、N/3個のパートができてしまいます。しかし、トータルでは、

$N/2 + N/3 + \cdots + N/M$ だから、 $O(N\log{M})$ 個のパートです。 $A_1 + A_2 + \cdots$ と $A_1 + 2A_2 + \cdots$ という二つの累積和を作っておけば $O(1)$ でそれぞれのパートが計算できます。

// You WILL Like Sigma Problem
#![allow(non_snake_case)]


//////////////////// library ////////////////////

fn read<T: std::str::FromStr>() -> T {
    let mut line = String::new();
    std::io::stdin().read_line(&mut line).ok();
    line.trim().parse().ok().unwrap()
}

fn read_vec<T: std::str::FromStr>() -> Vec<T> {
    read::<String>().split_whitespace()
            .map(|e| e.parse().ok().unwrap()).collect()
}


//////////////////// process ////////////////////

fn read_input() -> (Vec<i64>, Vec<i64>) {
    let _v: Vec<usize> = read_vec();
    let A: Vec<i64> = read_vec();
    let B: Vec<i64> = read_vec();
    (A, B)
}

const D: i64 = 998244353;

// -> [A1, A1+A2, ...]
fn accumulate(A: &Vec<i64>) -> Vec<i64> {
    let N = A.len();
    let mut v: Vec<i64> = vec![0; N+1];
    for i in 0..N {
        v[i+1] = (v[i] + A[i]) % D
    }
    v
}

// -> [A1, A1+2A2, ...]
fn accumulate2(A: &Vec<i64>) -> Vec<i64> {
    let N = A.len();
    let mut v: Vec<i64> = vec![0; N+1];
    for i in 0..N {
        v[i+1] = (v[i] + A[i] * (i as i64 + 1)) % D
    }
    v
}

fn F(A: Vec<i64>, B: Vec<i64>) -> i64 {
    let N = A.len();
    let M = B.len();
    let acc1 = accumulate(&A);
    let acc2 = accumulate2(&A);
    let mut counter: i64 = 0;
    for j in 2..M+1 {
        let mut c: i64 = 0;
        for first in (0..N).step_by(j) {
            // A[first]+2A[first]+...+(last-first)A[last-1]
            // = ((first-1)A[first]+...+last*A[last-1])
            //   - (first/j)(A[first]+...+A[last-1])
            let last = (first + j - 1).min(N);
            let q = (first / j * j) as i64;
            c = (c + acc2[last] - acc2[first] -
                    q * (acc1[last] - acc1[first])).rem_euclid(D)
        }
        counter = (counter + c * B[j-1]) % D
    }
    counter
}

fn main() {
    let (A, B) = read_input();
    println!("{}", F(A, B))
}