Module c
Expand description
ยงABC 369 C - Count Arithmetic Subarrays
refs: https://atcoder.jp/contests/abc369/tasks/abc369_c
use proconio::input;
fn main() {
input! {
n: usize,
a: [isize; n]
}
let diff = a.windows(2).map(|w| w[1] - w[0]).collect::<Vec<_>>();
let mut ans = n;
let run = run_length(&diff);
for (_, count) in run {
ans += (count * (count + 1)) / 2;
}
println!("{ans}");
}
fn run_length(s: &[isize]) -> Vec<(isize, usize)> {
s.iter().fold(Vec::new(), |mut acc, si| {
if let Some((last, count)) = acc.last_mut() {
if last == si {
*count += 1;
} else {
acc.push((*si, 1));
}
} else {
acc.push((*si, 1));
}
acc
})
}