AtCoder Beginner Contest 371 D

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

ソートして累積和を取って二分探索ですね。

// 1D Country
#![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<i32>, Vec<i64>, usize) {
    let _N: usize = read();
    let X: Vec<i32> = read_vec();
    let P: Vec<i64> = read_vec();
    let Q: usize = read();
    (X, P, Q)
}

fn read_query() -> (i32, i32) {
    let v: Vec<i32> = read_vec();
    let (L, R) = (v[0], v[1]);
    (L, R)
}

fn accumulate(v: Vec<(i32, i64)>) -> Vec<(i32, i64)> {
    let N = v.len();
    let mut a: Vec<(i32, i64)> = vec![(std::i32::MIN, 0); N+1];
    for (i, (x, p)) in v.into_iter().enumerate() {
        a[i+1] = (x, a[i].1 + p)
    }
    a
}

// a[i].0 >= xとなる最小のi
fn bin_search(first: usize, last: usize, x: i32, a: &Vec<(i32, i64)>) -> usize {
    if last - first <= 2 {
        if a[first].0 >= x {
            first
        } else if a[first+1].0 >= x {
            first + 1
        }
        else {
            last
        }
    }
    else {
        let mid = (first + last) / 2;
        if a[mid].0 >= x {
            if a[mid-1].0 < x {
                mid
            }
            else {
                bin_search(first, mid, x, a)
            }
        }
        else {
            bin_search(mid, last, x, a)
        }
    }
}

fn F_each(first: i32, last: i32, a: &Vec<(i32, i64)>) -> i64 {
    let i = bin_search(0, a.len(), first, a);
    let j = bin_search(0, a.len(), last, a);
    a[j-1].1 - a[i-1].1
}

fn F(X: Vec<i32>, P: Vec<i64>, Q: usize) {
    let mut villages: Vec<_> = X.into_iter().zip(P.into_iter()).collect();
    villages.sort();
    let acc = accumulate(villages);
    for _ in 0..Q {
        let (L, R) = read_query();
        println!("{}", F_each(L, R+1, &acc))
    }
}

fn main() {
    let (X, P, Q) = read_input();
    F(X, P, Q)
}