AtCoder Beginner Contest 429 D

https://atcoder.jp/contests/abc429/tasks/abc429_d

C以上になった人の数が何人なのかと、それが何回有効かを別にカウントするとコーディングが楽になります。前者は尺取り法的に計算すると計算量が減ります。

// On AtCoder Conference
#![allow(non_snake_case)]

use std::collections::HashMap;


//////////////////// 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() -> (usize, usize, Vec<usize>) {
    let v: Vec<usize> = read_vec();
    let (M, C) = (v[1], v[2]);
    let A: Vec<usize> = read_vec();
    (M, C, A)
}

// 場所ごとにカウントする
fn count(A: Vec<usize>) -> Vec<(usize, usize)> {
    let mut counter: HashMap<usize, usize> = HashMap::new();
    for a in A {
        let e = counter.entry(a).or_insert(0);
        *e += 1
    }
    let mut v: Vec<(usize, usize)> = counter.into_iter().collect();
    v.sort();
    v
}

fn collect_sums(v: &Vec<(usize, usize)>, C: usize) -> Vec<usize> {
    let L = v.len();
    let mut sums: Vec<usize> = vec![0; L];
    let mut last: usize = 0;
    let mut sum: usize = 0;
    for i in 0..L {
        while sum < C {
            sum += v[last].1;
            last = if last + 1 == L { 0 } else { last + 1 };
        }
        sums[i] = sum;
        sum -= v[i].1
    }
    sums
}

fn collect_spans(v: &Vec<(usize, usize)>, M: usize) -> Vec<usize> {
    let L = v.len();
    let mut spans: Vec<usize> = vec![0; L];
    spans[0] = M - v[L-1].0 + v[0].0;
    for i in 1..L {
        spans[i] = v[i].0 - v[i-1].0
    }
    spans
}

fn F(M: usize, C: usize, A: Vec<usize>) -> usize {
    let v = count(A);
    let sums = collect_sums(&v, C);
    let spans = collect_spans(&v, M);
    sums.into_iter().zip(spans.into_iter()).map(|(p, q)| p * q).sum::<usize>()
}

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