AtCoder Beginner Contest 371 E

https://atcoder.jp/contests/abc371/tasks/abc371_e

和を二重に取っている問題は和を取る順番を入れ替えるのが定番ですが、この二重和は対等なので、そこじゃないです。ある値を含む範囲がいくつあるかを数えます。ある値がp番目のみなら、かつだから、その範囲の個数はです。2個でp番目とq番目なら、などとなります。

// I Hate Sigma Problems
#![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()
}

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<i64> {
    let _N: usize = read();
    let A: Vec<i64> = read_vec();
    A
}

fn classify(A: Vec<i64>) -> Vec<Vec<usize>> {
    let N = A.len();
    let mut c: Vec<Vec<usize>> = vec![vec![]; N];
    for (p, a) in A.into_iter().enumerate() {
        c[a as usize - 1].push(p + 1)
    }
    c
}

fn F_each(ps: Vec<usize>, N: usize) -> usize {
    if ps.is_empty() {
        return 0
    }
    
    let mut counter: usize = ps[0] * (N - ps[0] + 1);
    for i in 1..ps.len() {
        counter += (ps[i] - ps[i-1]) * (N - ps[i] + 1)
    }
    counter
}

fn F(A: Vec<i64>) -> usize {
    let N = A.len();
    let c = classify(A);
    c.into_iter().map(|ps| F_each(ps, N)).sum::<usize>()
}

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