AtCoder Beginner Contest 318 E

https://atcoder.jp/contests/abc318/tasks/abc318_e

Aが同じ位置をまとめて、kが小さいほうから順に計算していけばよいです。

// Sandwiches
#![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() -> Vec<i32> {
    let _N: usize = read();
    let A = read_vec();
    A
}

fn collect_same_value_positions(A: &Vec<i32>) -> HashMap<i32, Vec<usize>> {
    let mut dic: HashMap<i32, Vec<usize>> = HashMap::new();
    for (i, &a) in A.iter().enumerate() {
        let e = dic.entry(a).or_insert(vec![]);
        e.push(i)
    }
    dic
}

fn g(ps: &Vec<usize>) -> usize {
    let h = |(n, s, a): (usize, usize, usize), d: usize| {
        let s1 = s + n * d + d - 1 - n;
        (n + 1, s1, a + s1)
    };
    
    let intervals = ps.iter().zip(&ps[1..]).map(|(p, q)| *q - *p);
    let (_, _, a) = intervals.fold((0, 0, 0), h);
    a
}

fn f(A: Vec<i32>) -> usize {
    let ps = collect_same_value_positions(&A);
    ps.values().map(g).sum::<usize>()
}

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