Module d

Expand description

§ABC 364 D - K-th Nearest

refs: https://atcoder.jp/contests/abc364/tasks/abc364_d

点 $B_j$ からの距離が $d$ 以内の中にある点 $A_i$ の個数が $K$ 以上となるような最小の $d$ を求めればよい. 最小の $d$ を求める二分探索と, 座標が $B_j - d$ 以上かつ $B_j + d$ 以下になる点 $A_i$ の個数を数える二分探索を組み合わせることで解くことができる. upper_bound, lower_bound が便利.

use proconio::input;
use superslice::Ext;

const INF: isize = 1e+9 as isize;

fn main() {
    input! {
        n: usize, q: usize,
        mut a: [isize; n],
        bk: [(isize, usize); q],
    }

    a.sort();

    for (b, k) in bk {
        if count_around(&a, b, 0) >= k {
            println!("0");
            continue;
        }

        let mut l = 0;
        let mut r = INF;

        while r - l > 1 {
            let m = (l + r) / 2;

            if count_around(&a, b, m) >= k {
                r = m;
            } else {
                l = m;
            }
        }

        println!("{r}");
    }
}

fn count_around(a: &[isize], base: isize, d: isize) -> usize {
    let upper_threshold = base + d;
    let lower_threshold = base - d;
    let lower_bound = a.lower_bound(&lower_threshold);
    let upper_bound = a.upper_bound(&upper_threshold);

    if upper_bound == a.len() - 1 && a[upper_bound] == upper_threshold {
        upper_bound - lower_bound + 1
    } else {
        upper_bound - lower_bound
    }
}