G - Pedometer 解説 by ngtkana


本解説では、添字を \(1\) ずらして \(A _ 0, \dots, A _ { n - 1 }\) にします。

数列 \(B = (B _ 0, \dots, B _ { n - 1 }, B _ n)\)\(B _ i = \sum _ { j \lt i } A _ j\) で定義すると答えは

\[ \left \lbrace (i, j) \middle | \begin{array}{l} 0 \le i < j \lt n, \\ B _ i = B _ j \pmod M \end{array} \right \rbrace + \left \lbrace (i, j) \middle | \begin{array}{l} 0 \le i < j \lt n, \\ B _ j = B _ n + B _ i \pmod M \end{array} \right \rbrace \]

に等しいので、あらかじめ \(A\) の総和を \(M\) で割ったあまりを求めておくことで、\(A\) を順に \(1\) 回走査して求めることができます。

use proconio::input;

fn main() {
    input! {
        n: usize,
        m: usize,
        a: [usize; n],
    }
    let sum = a.iter().sum::<usize>() % m;
    let mut prefix = 0;
    let mut count = vec![0; m];
    let mut ans = 0usize;
    for &x in &a {
        prefix = (prefix + x) % m;
        ans += count[prefix];
        ans += count[(prefix + m - sum) % m];
        count[prefix] += 1;
    }
    println!("{ans}");
}

提出: Rust 12 ms

投稿日時:
最終更新: