Zenn
👨‍🍳

ABC402:Rustで解く！問題解説

に公開
Rust
競技プログラミング
AtCoder
tech

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

A問題

問題

https://atcoder.jp/contests/abc402/tasks/abc402_a

解説

与えられた文字列 S からASCIIの大文字(A~Z)のみを抽出して、それらを連結した文字列を出力します。

コード

abc402a.rs
use proconio::{input, marker::Chars};
use itertools::Itertools;
fn main() {
    // 入力
    input!{
        s: Chars,
    }

    // 大文字のみ取得
    let mut ans = Vec::new();
    for &char in &s {
        if char.is_ascii_uppercase() { ans.push(char);}
    }

    // 出力
    println!("{}", ans.iter().join(""));
}

B問題

問題

https://atcoder.jp/contests/abc402/tasks/abc402_b

解説

キューを用いて以下の操作をシミュレーションします。

  1. クエリが 1 X の形式の場合、値 X をキューの末尾に追加します。
  2. クエリが 2 の形式の場合、キューの先頭の値を取り出して出力します。

コード

abc402b.rs
use proconio::input;
use std::collections::VecDeque;

fn main() {
    // クエリの数を入力
    input! {
        q: usize,
    }

    // キューを初期化
    let mut dq = VecDeque::new();

    // クエリ処理
    for _ in 0..q {
        // クエリの種類を入力
        input! {
            query_type: usize,
        }

        // クエリが「1 x」の場合
        if query_type == 1 {
            input! {
                x: usize,
            }
            dq.push_back(x);
        }
        // クエリが「2」の場合
        else if query_type == 2 {
            if let Some(val) = dq.pop_front() {
                println!("{}", val);
            }
        }
    }
}

C問題

問題

https://atcoder.jp/contests/abc402/tasks/abc402_c

解説

この問題では、N 種類の食材を順番に克服していくとき、i 日目に食べることができる料理の種類数を求めます。

各日に食べられる料理をそのまま調べると計算量が O(NM) となり、実行時間制限を超えてしまいます。そのため、効率的な方法を考える必要があります。

ここでは問題を「順番に食材を克服する」のではなく、「逆の順番に食材を嫌いになる」と考えます。これにより、嫌いになった食材が使われている料理は食べられなくなるという視点でシミュレーションを行います。

上記の視点でシミュレーションすると以下になります。

  1. 各食材が使用されている料理を事前にリスト化します。
  2. 現時点で食べられる料理の集合(HashSet)の要素数を記録します。
  3. 嫌いになる食材を後ろから順に処理し、その食材が使われている料理を集合から除外します。
  4. 各ステップで記録した食べられる要素数を逆順に出力します。

集合(HashSet)を使うことで、要素の追加・削除の計算量の平均 O(1) になり、各食材に関連する料理を効率的に管理できるため、計算量を削減できます。

コード

abc402c.rs
use proconio::{input, marker::Usize1};
use std::collections::HashSet;

fn main() {
    // 入力
    input! {
        n: usize, // 食材数
        m: usize, // 料理数
    }

    // 食材ごとに使用される料理のリストを作成
    let mut ingredients = vec![Vec::new(); n];
    for i in 0..m {
        input! {
            k: usize, // 料理に使われる食材の数
            a: [Usize1; k], // 料理に使われる食材のインデックス
        }
        for &ingredient in &a {
            ingredients[ingredient].push(i);
        }
    }

    // 克服する食材の順序
    input! {
        b: [Usize1; n],
    }

    // 結果を格納
    let mut ans = Vec::new();

    // 食べられる料理の集合
    let mut meal_set = HashSet::new();
    for i in 0..m {
        meal_set.insert(i);
    }

    // 後ろから食材を調べる
    for i in (0..n).rev() {
        // 現時点で食べられる料理の数を記録
        ans.push(meal_set.len());

        // 嫌いになった食材が使われている料理を除外
        for &meal in &ingredients[b[i]] {
            meal_set.remove(&meal);
        }
    }

    // 答えを逆順に出力
    for i in (0..n).rev() {
        println!("{}", ans[i]);
    }
}

D問題

問題

https://atcoder.jp/contests/abc402/tasks/abc402_d

解説

円周上に等間隔に存在する N 個の点について、2点を結んだ直線 M 本の交点の個数を求める問題です。
交点の個数を直接求めるのは難しいため、余事象を考えることで効率的に解を導きます。具体的な考え方としては以下のようになります。

  • 交点の個数は、直線を2本選ぶ組み合わせの総数 {}_M C_2 から、交点を作らないケースを引くことで求められます。
    • 交点を作らないケースとは、2本の直線が平行である場合です。
      • また、2本の直線が平行であるとは、円周上の点 (a, b) を結ぶ直線について、(a + b) \mod N が同じ場合に成り立ちます。

そのため、各直線について (a + b) \mod N の値を計算し、その値ごとに直線のパターンを分類します。同じパターンの直線から2本を選ぶ組み合わせの総数を計算し、これを {}_M C_2 から引きます。

これにより、交点の個数を効率的に求めることができます。

コード

abc402d.rs
use proconio::{input, marker::Usize1};

fn main() {
    // 入力
    input! {
        n: usize, // 円周上の点の数
        m: usize, // 直線の本数
        ab: [(Usize1, Usize1); m], // 直線を結ぶ点のペア
    }

    // 交点の個数 = 全ての直線2本の組み合わせ(MC2) - 平行な直線の組み合わせ
    let mut ans = m * (m - 1) / 2;

    // 各直線のパターン「(a + b) % n」を調べる
    let pattern = count_line_patterns(&ab, n);

    // 同じパターンから2本選ぶ組み合わせを引く
    ans -= calc_same_pattern(&pattern);

    // 出力
    println!("{}", ans);
}

// 各直線の「(a + b) % n」のパターンをカウントする
fn count_line_patterns(lines: &Vec<(usize, usize)>, n: usize) -> Vec<usize> {
    let mut pattern = vec![0; n];
    for &(from, to) in lines {
        let p = (from + to) % n;
        pattern[p] += 1;
    }
    pattern
}

// 同じパターンから2本選ぶ組み合わせを計算する
fn calc_same_pattern(pattern: &Vec<usize>) -> usize {
    let mut ret = 0;
    for &p in pattern {
        if p > 1 {
            ret += p * (p - 1) / 2;
        }
    }
    ret
}

E問題

問題

https://atcoder.jp/contests/abc402/tasks/abc402_e

解説

この問題は、N個の問題が与えられ、それぞれの問題にはスコア、解答に必要なコスト、正解する確率が設定されています。与えられた所持金の範囲内で問題を解く際の総スコアの期待値を最大化します。

この問題は動的計画法（DP）を用いて解くことができます。DPの状態を次のように定義します。

  • DP[\text{所持金}][\text{問題の正解状況}] = これから得られるスコアの期待値の最大値

DPの遷移は以下のように行います。

  1. 現在の所持金と正解状況から、次に解く問題を選択します。

  2. その問題を解くことで得られるスコアの期待値を計算します。この期待値は以下の式で求められます。

    \begin{aligned} \text{問題i選択時のスコア期待値} &= (\text{現在のスコア} + \text{問題のスコア}) \times \text{正解確率} \\ &= (\text{現在のスコア}) \times (\text{1 - 正解確率}) \\ \end{aligned}

  3. 2.について、その時点で選択できる全ての問題に対して期待値を計算し、期待値が最も大きくなる値でDPテーブルを更新します。

全体の計算量は O(NX2^N) になりますが、N の制約が 8 以下と小さいため、実行時間制限内で解くことが可能です。

コード

abc402e.rs
use proconio::input;
fn main() {
    // 入力
    input!{
        n: usize, x: usize,  // 問題数、所持金
        scp: [(usize, usize, usize); n], // (スコア, コスト, 正解確率) のタプル
    }
    
    // DP[所持金][正解状態] = これから得られるスコアの期待値の最大値
    let mut dp = vec![vec![0.0; 1 << n]; x+1];

    // もらうDP。所持金、正解状態を全て調べる。
    for cur_yen in 0..=x {
        for cur_state in 0..1<<n {
            let mut max_score = 0.;

            // 選択する問題を全探索
            for bits in 0..n {
                // 正解済みの問題を選択した場合、お金が足りない場合は見ない
                if cur_state & (1 << bits) != 0 { continue; }
                if scp[bits].1 > cur_yen { continue; }

                // 操作後の所持金、正解状態
                let next_yen = cur_yen - scp[bits].1;
                let next_state = cur_state | (1 << bits);
                
                // スコアの期待値を計算
                let score = calc_score(scp[bits].0, scp[bits].2 as f64 / 100., 
                    &dp, cur_state, next_state, next_yen);
                max_score = fmax(&vec![max_score, score]);
            }
            
            // スコアの期待値が最も大きくなるものを採用して、更新
            dp[cur_yen][cur_state] = fmax(&vec![dp[cur_yen][cur_state], max_score]);
        }
    }
    
    // 出力
    println!("{}", dp[x][0]);

}

// スコアの期待値を計算する関数
fn calc_score(score: usize, px : f64, 
    dp: &Vec<Vec<f64>>, cur_state: usize, next_state: usize, next_yen: usize) -> f64{
    let mut ret = 0.;

    // スコアを得た場合
    ret += (dp[next_yen][next_state] + score as f64) * px;
    // スコアを得られなかった場合
    ret += (dp[next_yen][cur_state]) * (1. - px);

    ret
}

fn fmax(x: &Vec<f64>) -> f64 {
    let mut ret = x[0];
    for &i in x {
        if ret < i {
            ret = i;
        }
    }
    ret
}

Discussion