AtCoder Beginner Contest 440 D

https://atcoder.jp/contests/abc440/tasks/abc440_d

XからX+YがソートしたAの要素が何個またぐかを二分探索します。

// Forbidden List 2
#![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() -> (usize, Vec<usize>) {
    let v: Vec<usize> = read_vec();
    let Q = v[1];
    let A: Vec<usize> = read_vec();
    (Q, A)
}

fn read_query() -> (usize, usize) {
    let v: Vec<usize> = read_vec();
    let (X, Y) = (v[0], v[1] as usize);
    (X, Y)
}

// eより小さい最大の要素のindex
fn find_index(e: usize, A: &Vec<usize>) -> usize {
    if A[0] >= e {
        return usize::MAX
    }
    
    let mut first: usize = 0;
    let mut last: usize = A.len();
    while last - first > 1 {
        let mid = (first + last) / 2;
        if A[mid] < e {
            first = mid
        }
        else {
            last = mid
        }
    }
    first
}

fn F_each(X: usize, Y: usize, A: &Vec<usize>) -> usize {
    let N = A.len();
    let i_ = find_index(X, A);
    // 最初にまたぐindex
    let i0 = if i_ == usize::MAX { 0 } else { i_+1 };
    if i0 == N {
        return X + Y - 1
    }
    else if X + Y <= A[i0] {
        return X + Y - 1
    }
    
    // Y + X + (i - i0) = A[i] + r (r >= 0)
    // となるような最大のiを求める
    let mut first = i0;
    let mut last = A.len();
    while last - first > 2 {
        let mid = (first + last) / 2;
        if Y + mid - i0 + X <= A[mid] {
            last = mid
        }
        else {
            first = mid
        }
    }
    if first < last - 1 && Y + first + 1 - i0 + X > A[first+1] {
        first += 1
    }
    Y + first - i0 + X
}

fn F(Q: usize, mut A: Vec<usize>) {
    A.sort();
    for _ in 0..Q {
        let (X, Y) = read_query();
        println!("{}", F_each(X, Y, &A))
    }
}

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