AtCoder Beginner Contest 310 E

https://atcoder.jp/contests/abc310/tasks/abc310_e

DPですね。でを固定したとき、0になるものと1になるもののカウントしたものを状態と考えて、の状態を計算します。

// NAND repeatedly
#![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()
}


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

fn read_input() -> String {
    let _N: usize = read();
    let S = read();
    S
}

fn f(S: String) -> u64 {
    let v: Vec<i32> = S.chars().map(|c| (c as i32) - ('0' as i32)).collect();
    let mut dp: (u64, u64) = (0, 0);
    let mut total: u64 = 0;
    for b in v.into_iter() {
        if b == 0 {
            dp = (1, dp.0 + dp.1)
        }
        else {
            dp = (dp.1, dp.0 + 1)
        }
        total += dp.1
    }
    total
}


//////////////////// main ////////////////////

fn main() {
    let S = read_input();
    println!("{}", f(S))
}