👨‍🍳

ABC390:Rustで解く!問題解説

2025/01/26に公開

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

A問題

  • 問題

https://atcoder.jp/contests/abc390/tasks/abc390_a

  • 解説

整数列 A について、隣接する2つの要素を実際に交換して、昇順に並び替えたものと一致するかどうかをチェックします。
一致した場合はその時点で Yes と出力し、違った場合は元の状態に戻します。この操作を全ての隣接する要素のペアに対して試みます。

  • コード
use proconio::input;

fn main() {
    input!{
        mut a: [usize; 5],
    }

    let correct_a = vec![1, 2, 3, 4, 5];

    for i in 0..4 {
        // 隣接する2つを交換する
        a.swap(i, i + 1);

        // 昇順のものと一致したらOK
        if a == correct_a {
            println!("Yes");
            return;
        }

        // 一致しないなら戻す
        a.swap(i, i + 1);
    }
    
    println!("No");
}

B問題

  • 問題

https://atcoder.jp/contests/abc390/tasks/abc390_b

  • 解説

整数列 A が等比数列であるためには、以下の形を満たす必要があります。

A_1, rA_1, r^2A_1, \ldots, r^{N-1}A_1

この状態で両側から見ていき2つの値の積を求めると、以下の通り、積の値が全て一致するので、それを確認すればよいです。

  • 左から 1 番目と右から 1 番目の積: A_1 \times r^{N-1}A_1 = r^{N-1}A_1^2

  • 左から 2 番目と右から 2 番目の積: rA_1 \times r^{N-2}A_1 = r^{N-1}A_1^2
    ...

  • 左から N 番目と右から N 番目の積: A_1 \times r^{N-1}A_1 = r^{N-1}A_1^2

  • コード

use proconio::input;

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

    // 両側から2つの積を取る
    let val = a[0] * a[n-1];
    for i in 0..n {
        // 左からi番目と右からi番目の積がvalと一致するかを確認
        if a[i] * a[n-1-i] != val {
            println!("No");
            return;
        }
    }
    println!("Yes");
}

C問題

  • 問題

https://atcoder.jp/contests/abc390/tasks/abc390_c

  • 解説

マス目の中にある黒マス(#)をチェックし、①「縦座標の最小値・最大値、横座標の最小値・最大値」を調べます。
①の範囲にあるマスに白マス(.)が含まれていなければ、Yes を出力し、そうでなければ No を出力します。

  • コード
use proconio::{input, marker::Chars};
use std::cmp::{min, max};

fn main() {
    input!{
        h: usize, w: usize,
        s: [Chars; h],
    }

    // #マスの縦のmin-max, 横のmin-maxを取得
    let mut row_min_max = (h, 0);
    let mut col_min_max = (w, 0);
    for i in 0..h {
        for j in 0..w {
            if s[i][j] == '#' {
                // 縦のminを更新
                row_min_max.0 = min(row_min_max.0, i);
                // 縦のmaxを更新
                row_min_max.1 = max(row_min_max.1, i);
                // 横のminを更新
                col_min_max.0 = min(col_min_max.0, j);
                // 横のmaxを更新
                col_min_max.1 = max(col_min_max.1, j);
            }
        }
    }

    // [min..=max]の範囲に'.'がないかをチェック
    for i in row_min_max.0 ..= row_min_max.1 {
        for j in col_min_max.0 ..= col_min_max.1 {
            if s[i][j] == '.' {
                println!("No");
                return;
            }
        }
    }
    println!("Yes");
}

D問題

  • 問題

https://atcoder.jp/contests/abc390/tasks/abc390_d

  • 解説

N 個の要素からなる配列 A を1つ以上のグループに割り振ります。それぞれの要素について、①「既にあるグループのどこに含めるか」または②「新しいグループに追加するか」を再帰関数を用いて順番に調べていきます。
N 個全ての要素を選んだ状態でXORを計算し、その結果得られた候補を集合に追加して、候補数を最終的な答えとします。

  • コード
use proconio::input;
use std::collections::HashSet;

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

    // 答えの種類数を格納する集合
    let mut ans_st = HashSet::new();
    // 各グループの値を保持するベクタ
    let mut groups = Vec::new();
    // 再帰関数を実行
    rec(0, n, &a, &mut groups, &mut ans_st);

    // 結果として、集めたXORの値の種類数を出力
    println!("{}", ans_st.len());
}

fn rec(
    cnt: usize, n: usize, // 見た個数, 全体の個数
    a: &Vec<usize>, groups: &mut Vec<usize>, // Aの候補, 各グループの値
    ans_st: &mut HashSet<usize> // 答えの種類数
) {
    // 全ての要素を見た場合
    if cnt == n {
        // 現在のグループのXOR計算
        let val = calc_xor(groups);
        // 結果を集合に追加
        ans_st.insert(val);
        return;
    }
    
    // 既にあるグループに要素を追加する、または新しいグループを作成する処理

    // 既にあるグループを選択
    let sz = groups.len();
    for i in 0..sz {
        groups[i] += a[cnt];
        rec(cnt + 1, n, a, groups, ans_st);
        groups[i] -= a[cnt];
    }

    // 新しいグループを追加
    groups.push(a[cnt]);
    rec(cnt + 1, n, a, groups, ans_st);
    groups.pop(); // 最後に追加したグループを削除
}

// 全グループのXORを求める関数
fn calc_xor(vec: &Vec<usize>) -> usize {
    let mut ret = 0;
    for i in 0..vec.len() {
        ret ^= vec[i]; // 各グループの値をXOR計算
    }
    ret
}

E問題

  • 問題

https://atcoder.jp/contests/abc390/tasks/abc390_e

  • 解説

食べ物1つにつきビタミンは1つのみとなっているので、ビタミンの種類毎に選んだ結果の組み合わせで解くことができます。ビタミンの種類毎の選び方はDP(動的計画法)を用いることで、
カロリー j に対する摂取量の最大値を求めます。これをビタミン1~3で事前に計算します。

次に、ビタミン1〜3の組み合わせ方ですが、全てのパターンを試そうとすると O(X^3) となり、実行時間内に解くことができないため、効率化を図ります。具体的には、答えとなる摂取量 a の値を二分探索で求める方法で解きます。
ビタミン1〜3で摂取量 a' を固定したときに、①「摂取量 a' 以上となるカロリー b 」を見つけて、カロリーの合計が X 以下にできるかどうかを判定して絞り込むことができます。
なお、①を見つける過程でも二分探索が必要で、「カロリー b 以下で摂取量が最も大きい値」を見つけるために、摂取量の累積の最大値を事前に調べておく必要があります。

  • コード
use proconio::{input, marker::Usize1};
use std::cmp::max;
use std::mem::swap;
const INF: i64 = 1 << 60;

fn main() {
    input!{
        // 食べ物の個数, カロリー上限
        n: usize, x: usize,
    }
    
    // vの種類毎に、データをまとめる
    let mut data = vec![Vec::new(); 3];
    for _ in 0..n {
        input!{
            v: Usize1, a: usize, c: usize,
        }
        // ビタミン毎の(摂取量, カロリー)
        data[v].push((a, c));
    }

    // v1 ~ v3をそれぞれDP
    // dp[現在カロリーj] = 摂取量のmax
    let mut dp1 = vec![-INF; x+1]; dp1[0]=0;
    let mut dp2 = vec![-INF; x+1]; dp2[0]=0;
    let mut dp3 = vec![-INF; x+1]; dp3[0]=0;
    calc_dp(&mut dp1, &data[0], x);
    calc_dp(&mut dp2, &data[1], x);
    calc_dp(&mut dp3, &data[2], x);

    // v1 ~ v3のDPテーブルを累積maxにする。
    calc_acc_max(&mut dp1, x);
    calc_acc_max(&mut dp2, x);
    calc_acc_max(&mut dp3, x);

    // 答えを二分探索
    let ans = calc_binary_search(&dp1, &dp2, &dp3, x);
    println!("{}", ans);

}

fn calc_dp(now_dp: &mut Vec<i64>, data: &Vec<(usize, usize)>, x: usize) {
    let sz = data.len();
    // もらうDP
    for i in 1..=sz {
        let mut next_dp = vec![-INF; x+1];
        for j in 0..=x {
            // 選択しない場合
            next_dp[j] = max(next_dp[j], now_dp[j]);
            // 選択する場合
            if j >= data[i-1].1 {
                next_dp[j] = max(next_dp[j], now_dp[j - data[i-1].1] + data[i-1].0 as i64);
            }
        }
        swap(now_dp, &mut next_dp);
    }
}

fn calc_acc_max(dp: &mut Vec<i64>, x: usize) {
    for i in 0..x {
        dp[i+1] = max(dp[i+1], dp[i]);
    }
}

fn calc_binary_search(dp1: &Vec<i64>, dp2: &Vec<i64>, dp3: &Vec<i64>, x: usize) -> usize {
    let mut ok = 0;
    let mut ng = INF;
    for _ in 0..100 {
        // 摂取量が全てmid以上になるように選択して、カロリーがx以下かどうかで絞り込み
        let mid = (ok + ng) / 2;
        let pos1 = lower_bound(&dp1, mid);
        let pos2 = lower_bound(&dp2, mid);
        let pos3 = lower_bound(&dp3, mid);
        if pos1 + pos2 + pos3 > x { ng = mid;}
        else { ok = mid;}
    }
    ok as usize
}

// 二分探索
// x以上の最小のposを返す 
pub fn lower_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;
        }
    }
    l
}

Discussion