👨‍🍳

ABC382:Rustで解く!問題解説

2024/12/01に公開

AtCoder Beginner Contest 382のA~E問題をRustで解いた際の解法をまとめました。

A問題

  • 問題

https://atcoder.jp/contests/abc382/tasks/abc382_a

  • 解説

与えられた文字列Sに含まれる@.にD個置き換えます。
よって、もともとの.の数に、Dを加えた値が答えとなります。

  • コード
use proconio::{input, marker::Chars};

fn main() {
    input! {
        _n: usize, // 使用しない変数_n
        d: usize,  // '@'を'.'に置き換える個数
        s: Chars,  // 文字列 S
    }

    let mut ans = 0; // もともと存在する '.' の数をカウントする変数
    for &char in &s {
        if char == '.' {
            ans += 1; // '.' の数をカウント
        }
    }
    ans += d; // '@' を '.' に置き換えた数を加える
    println!("{}", ans); // 最終的な結果を出力
}

B問題

  • 問題

https://atcoder.jp/contests/abc382/tasks/abc382_b

  • 解説

与えられた文字列Sに含まれる@を右側からD.に置き換えます。
置き換えが完了したら、変更後の文字列Sを出力します。

  • コード
use proconio::{input, marker::Chars};
use itertools::Itertools;

fn main() {
    input! {
        n: usize,  // 文字列の長さ
        d: usize,  // 置き換える '@' の数
        mut s: Chars,  // 文字列 S を文字のベクタとして受け取る
    }

    let mut cnt = 0;  // 置き換えた '@' の数をカウントする
    for i in (0..n).rev() {  // 右側から左に向かって走査
        if cnt == d { break; }  // すでに D 個置き換えたらループを終了
        if s[i] == '@' {  // '@' に出会った場合
            s[i] = '.';  // '@' を '.' に置き換える
            cnt += 1;  // 置き換えた数をインクリメント
        }
    }
    // 置き換え後の文字列を出力
    println!("{}", s.iter().join(""));
}

C問題

  • 問題

https://atcoder.jp/contests/abc382/tasks/abc382_c

  • 解説

そのまま実装すると、寿司の数Mと食べる人の数Nにより、計算量はO(NM)となるため、実行時間内に間に合いません。
そこで、食べられない人の情報を捨てることを考えます。
具体的には、数列Aを左から順に見た際に、その時点の最小値以上の要素は、前の人に食べられてしまうため、不要な情報となります。不要な情報を削除した数列をA'とすると、A'内の各要素A_i以上の値を持つ寿司は、その人が全て食べるとみなすことができます。
この状態で数列 Bの一つひとつの要素について、二分探索を用いてA' からB_i以上のA_iを見つけ、その人のインデックスを出力します。

  • コード
use proconio::input;
use std::collections::HashMap;
const INF: usize = 1 << 60;

fn main() {
    input! {
        n: usize, m: usize,
        a: [usize; n],
        b: [usize; m],
    }

    // aを順に見て、現在の最小値より小さい要素だけを抽出する
    let mut aa = Vec::new(); // aの抽出データ(昇順)
    let mut a_map = HashMap::new(); // aの抽出データの紐づけ情報
    let mut a_min = INF; // aの最小値

    // Aから不要な要素を取り除く
    for i in 0..n {
        if a[i] < a_min {
            // 新しい最小値を記録
            a_map.insert(a[i], i + 1); // 1-based indexを保存
            aa.push(a[i]);
            a_min = a[i];
        }
    }
    aa.sort();

    // Bの各要素に対して二分探索を実行
    for i in 0..m {
        let pos = below_bound(&aa, b[i]);
        if pos == INF {
            println!("-1");
        } else {
            let val = aa[pos];
            println!("{}", *a_map.get(&val).unwrap());
        }
    }
}

// 二分探索
// x以下の最大のposを返す。0番目より小さい値では異常値として、INFを返す
pub fn below_bound<T: std::cmp::PartialOrd>(vec: &Vec<T>, val: T) -> usize {
    let mut l = 0;
    let mut r = vec.len();
    while r - l > 0 {
        let m = (l + r) / 2;
        if vec[m] <= val {
            l = m + 1;
        } else {
            r = m;
        }
    }
    if l == 0 { INF } else { l - 1 }
}

D問題

  • 問題

https://atcoder.jp/contests/abc382/tasks/abc382_d

  • 解説

問題文と入力例1、出力例1の内容から以下のことが分かります。

  • N個の数列を作成します。
  • 初項は1以上。
  • 末項はM以下。
  • 数列の2項の間は10以上の差があります。

最小ケース(10N - 9)の場合、数列は1, 11, ..., 10(N-1) + 1の1つのみになります。
最小ケースでない場合は、次の式で計算される値だけ余分に足すことができます。

\begin{aligned} M - (10(N - 1) + 1) &= M - 10N + 10 - 1 \\ &= M + 9 - 10N \end{aligned}

この値分だけ、N項の中で余分に加算できるため、後いくつ追加で加算できるかを保持し、再帰関数を使って全パターンを列挙します。

  • コード
use proconio::input;
use itertools::Itertools;

fn main() {
    input! {
        n: usize, m: usize,
    }

    let mut ans = vec![];
    let mut vec = vec![];
    
    // 追加可能な値を計算
    let rest_add = m + 9 - 10 * n;

	// 再帰関数
    rec(0, &mut vec, rest_add, &mut ans, n);

    // 選択数, 選択した値リスト, 残りの加算値, 必要な選択個数
    fn rec(cnt: usize, vec: &mut Vec<usize>, rest_add: usize, ans: &mut Vec<Vec<usize>>, n: usize) {
        // n個選択したら結果に追加
        if vec.len() == n {
            ans.push(vec.clone());
            return;
        }

        for i in 0..=rest_add {
            let mut val = 1; // 初項の設定
            if !vec.is_empty() {
                // 以前の項に10を加算
                val = vec[vec.len() - 1] + 10;
            }
            vec.push(val + i); // 現在の選択値を追加
            rec(cnt + 1, vec, rest_add - i, ans, n); // 再帰呼び出し
            vec.pop(); // 最後に追加した値を除去
        }
    }

    // 結果の出力
    println!("{}", ans.len()); // 結果の数を出力
    for item in &ans {
        println!("{}", item.iter().join(" ")); // 各結果を出力
    }
}

E問題

  • 問題

https://atcoder.jp/contests/abc382/tasks/abc382_e

  • 解説

この問題は2つの動的計画法 (DP) を使って解くことができます。

(1) 1パックごとに手に入るレアの枚数の各確率を求めるDP
dp[i枚目まで見た時][入手したレアの枚数j] = 確率
上記はi枚目を取得した際に、そのカードがレアかそうでないかによるDPの遷移を行います。

(2) 期待値を求めるDP
dp[レアがi枚の状態になる] = パックを引く回数の期待値
これは、パックを引く前に持っていたレアの枚数からの遷移を考えることで求められます。
1パックが2枚なので、1回引いて0~2枚もらいレアがi枚の状態にする。
ただし、0枚もらうというのは自身を指すので、その状態が発生した場合はやり直しが必要になります。
0枚もらう確率をp_0とすると、正しく遷移する確率は(1 - p_0)となるので、その確率で割る必要があります。

具体例として、入力例2において、レア2枚の状態の期待値を求める場合、以下のような計算になります。

\begin{aligned} dp[2] &= (dp[0] * p_2 + dp[1] * p_1 + dp[2] * p_0 + 1) \\ (1 - p_0) * dp[2] &= (dp[0] * p_2 + dp[1] * p_1 + 1) \\ dp[2] &= (dp[0] * p_2 + dp[1] * p_1 + 1) / (1 - p_0) \end{aligned}
  • コード
use proconio::input;

fn main() {
    input!{
        n: usize, x: usize,
        mut p: [f64; n],
    }
    
    // pの確率(%)を小数に変換する
    for i in 0..n {
        p[i] /= 100.0;
    }

    // 1パックごとに手に入るレアの枚数の各確率を求めるDP
    // dp[i枚目まで見た時][入手したレアの枚数j] = 確率
    let mut dp_px = vec![vec![0.0; n + 1]; n + 1];
    dp_px[0][0] = 1.0;

    // 確率DP
    for i in 1..=n {
        // カードiを受け取って、当たりがj枚になる状態
        for j in (0..=n).rev() {
            // カードi枚目が当たり
            if j != 0 {
                dp_px[i][j] += dp_px[i - 1][j - 1] * p[i - 1];
            }
            // カードi枚目が外れ
            dp_px[i][j] += dp_px[i - 1][j] * (1. - p[i - 1]);
        }
    }
    // dp_px[n]が、1パックごとに手に入るレアの枚数の各確率

    // 期待値DP
    // dp[レアがi枚の状態になる] = パックを引く回数の期待値
    let mut dp = vec![0.0; x + 1];
    for i in 1..=x {
        // 期待値計算
        let mut tot = 0.0;
        for j in 1..=n {
            if i >= j {
                tot += dp_px[n][j] * dp[i - j];
            }
        }
        tot += 1.0;
        // 正しく遷移する確率で割る
        if dp_px[n][0] != 0. {
            tot /= 1. - dp_px[n][0];
        }
        dp[i] = tot;
    }
    println!("{}", dp[x]);
}

Discussion