AtCoder Beginner Contest 321 D

https://atcoder.jp/contests/abc321/tasks/abc321_d

Aは昇順、Bは降順にして、Bの累積を計算しておいて、しゃくとり法っぽく計算します。

// Set Menu
#![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() -> (i64, Vec<i64>, Vec<i64>) {
    let v = read_vec();
    let P: i64 = v[2];
    let A = read_vec();
    let B = read_vec();
    (P, A, B)
}

fn accumulate(B: &Vec<i64>) -> Vec<i64> {
    let mut v: Vec<i64> = vec![0; B.len()+1];
    for (i, b) in B.iter().enumerate() {
        v[i+1] = v[i] + b
    }
    v
}

fn F(P: i64, mut A: Vec<i64>, mut B: Vec<i64>) -> i64 {
    A.sort();
    B.sort_by(|a, b| b.cmp(a));     // 降順
    let acc_B = accumulate(&B);
    let L = B.len();
    let mut s: i64 = 0;
    let mut j: usize = 0;
    for a in A.into_iter() {
        while j < L && a + B[j] > P {
            j += 1
        }
        s += P * (j as i64) + a * ((L-j) as i64) + acc_B[L] - acc_B[j]
    }
    s
}

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