AtCoder Beginner Contest 393 D

https://atcoder.jp/contests/abc393/tasks/abc393_d

要するに1の移動距離の和を求めればよいですが、1の位置の累積和を計算しておけばどこを中心に集まるかを決めるとO(1)で求められるので、全ての中心について計算すればよいです。

// Fine Triplets
#![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: String = read();
    S
}

fn T(n: usize) -> usize {
    if n == 0 { 0 } else { n * (n - 1) / 2 }
}

fn F(S: String) -> usize {
    let N = S.len();
    let ps: Vec<usize> = S.chars().enumerate().
                        filter(|&(_, c)| c == '1').map(|(p, _)| p).collect();
    let M = ps.len();
    let mut acc1: Vec<usize> = vec![0; M+1];
    for i in 1..M+1 {
        acc1[i] = acc1[i-1] + ps[i-1]
    }
    let mut acc2: Vec<usize> = vec![0; M+1];
    for i in (0..M).rev() {
        acc2[i] = acc2[i+1] + ps[i]
    }
    
    let mut min_swaps = N * N;
    let mut k: usize = 0;
    let mut l: usize = 0;
    for p in 0..N {
        if l < M && p == ps[l] {
            l += 1
        }
        let s1 = (p + 1 - l) * l + T(l) - acc1[l];
        let s2 = acc2[l] + T(l) - T(M) - (p - k) * (M - l);
        let s = s1 + s2;
        if s < min_swaps {
            min_swaps = s
        }
        if k < M && p == ps[k] {
            k += 1
        }
    }
    min_swaps
}

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