Rustコトハジメ

プログラミング言語Rustと競プロに関する情報をお届けします。

【yukicoder No.871 かえるのうた】配列を埋めるためにDFSする教育的良問

www.rustforbeginners.com

で解けるべきだった問題のやり直しACです。

二分探索の使い方と、配列の上でDFSするという組み合わせの、教育的な良問だと思いました。

No.871 かえるのうた - yukicoder

問題

Xiにいるかえるが鳴くと、[Xi-Ai,Xi+Ai]にいるかえるが鳴く。これが伝搬していく。K番目のかえるが最初に鳴くとき、最終的に何匹のかえるが鳴くか?

方針

i番目のかえるが鳴いた時、どのかえるが鳴くかは二分探索で計算出来る。

その上で、そいつらで鳴くことを深さ優先的に実行していく。(iが鳴いて、jが鳴くならば、iからjに枝があるのと同じこと)

愚直にループして配列を埋めていくのではなく、DFSするところがかっこいいところ。

実装

Rustのネスト関数は環境にアクセスしないし、クロージャ再帰出来ないため、Rustでの再帰は本当にいや。

fn dfs(i: usize, b: &mut [bool], X: &[i64], A: &[i64]) {
    b[i]=true;
    let lower = X[i]-A[i];
    let upper = X[i]+A[i];
    let k0 = X.lower_bound(&lower);
    let k1 = X.lower_bound(&(upper+1));
    for j in k0..k1 {
        if b[j] { continue; }
        dfs(j, b, X, A);
    }
}
fn solve() {
    input!{
        N:usize,K:usize,
        X:[i64;N],
        A:[i64;N],
    }
    let mut b=vec![false;N];
    dfs(K-1, &mut b, &X, &A);

    let mut n=0;
    for i in 0..N{
        if b[i] {
            n+=1;
        }
    }
    println!("{}",n);
}

考察

直接的にはグラフではないけど、グラフのようなものを想像して、それに対してDFSするというのはよくある。メモ化再帰を考えれば、ある意味自然。