Zenn
👨‍🍳

ABC396:Rustで解く!問題解説

2025/03/09に公開

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

A問題

問題

https://atcoder.jp/contests/abc396/tasks/abc396_a

解説

隣接する3つの値が同じかどうかをチェックします。全て試して1箇所でも見つかればYes 、見つからなければ No を出力します。

コード

abc396a.rs
use proconio::input;
fn main() {
    // 入力
    input!{
        n: usize,
        a: [usize; n],
    }

    let mut f = false;
    for i in 0..n-2 {
        if a[i] == a[i+1] && a[i+1] == a[i+2] { f = true;}
    }
    
    // 出力
    println!("{}", if f {"Yes"} else {"No"});
}

B問題

問題

https://atcoder.jp/contests/abc396/tasks/abc396_b

解説

スタックのデータ構造を用いて解くことができます。問題文の通り、0を100個ストックした状態で始めて、クエリが1ならスタックの一番上に要素を追加し、クエリが2ならスタックの一番上の要素を取り出して出力します。

コード

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

fn main() {
    // スタックを初期化
    let mut dq = VecDeque::new();
    for _ in 0..100 { dq.push_back(0);}

    // 入力
    input!{
        q: usize,
    }

    // クエリ処理
    for _ in 0..q {
        input!{
            nm: usize,
        }
        // 1なら、追加
        if nm == 1 {
            input!{
                x: usize,
            }
            dq.push_back(x);
        }
        // 2なら、スタックの先頭を取り出して出力
        else {
            let x = dq.pop_back().unwrap();
            println!("{}", x);
        }
    }
}

C問題

問題

https://atcoder.jp/contests/abc396/tasks/abc396_c

解説

方針として、価値が高いボールから優先的に取ることで総価値を最大化できます。ボールの取り方は上記を前提とし、ボールの価値は降順に並び替えておきます。
次に、ボールの個数の決め方ですが、黒と白を両方全探索すると計算量が O(NM)O(NM) となり、実行制限時間内に解くことができません。そのため、工夫が必要です。
ここでは、黒と白のボールを同じ個数取る全探索ができるようにします。ボールの個数は黒が白以上である必要があるため、黒を基準に考えた場合、無価値な白(個数が足りていない、または価値が0未満)のボールは取らないようにします。具体的には、無価値な白のボールは全て0に置き換え、価値0のダミーボールを扱うことにします。
これにより、黒のボールの個数を全探索するだけで、白のボールの取り方も適切に選択でき、計算量が O(N)O(N) に減少するため、実行制限時間内に解くことが可能になります。

コード

abc396c.rs
use proconio::input;
use std::cmp::max;

fn main() {
    // 入力
    input!{
        n: usize, mut m: usize,
        mut b: [i64; n],
        mut w: [i64; m],
    }

    // wの価値がマイナスのものを0にする
    white_more0(&mut w, m);
    // wの個数がbと同じ個数以上になるように0を追加
    white_add0(&mut w, &mut m, n);
    // b,wの価値を降順に並べ替える
    b.sort(); b.reverse();
    w.sort(); w.reverse();

    // bの選ぶ個数を全探索し、wも同じ個数選んで、価値のmaxを計算する
    let mut ans = 0;
    let mut cur = 0;
    for i in 0..n {
        cur += b[i] + w[i];
        ans = max(ans, cur);
    }

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

// wの各要素が0未満の場合に0に置き換える関数
fn white_more0(w: &mut Vec<i64>, m: usize) {
    for i in 0..m {
        if w[i] < 0 { 
            w[i] = 0; 
        }
    }
}

// wの要素数がnに満たない場合に0を追加する関数
fn white_add0(w: &mut Vec<i64>, m: &mut usize, n: usize) {
    while n > *m {
        w.push(0);
        *m += 1;
    }
}

D問題

問題

https://atcoder.jp/contests/abc396/tasks/abc396_d

解説

DFS(深さ優先探索)を用いて、頂点 11 から 頂点 NN への経路を全探索することで解くことができます。探索時に通過した辺の値のXORを取っていき、頂点 NN に到達した際に総XORの最小値を更新すればよいです。

コード

abc396d.rs
use proconio::{input, marker::Usize1};
use std::cmp::min;
const INF: usize = 1 << 60;
fn main() {
    // 入力
    input!{
        n: usize, m: usize,
        uvw: [(Usize1, Usize1, usize); m] // インデックスを0始まりにする
    }

    // 無向グラフの隣接リスト(向かいの頂点、辺の値)
    let mut edge = vec![Vec::new(); n];
    for (u, v, w) in uvw {
        edge[u].push((v, w));
        edge[v].push((u, w));
    } 

    // 答えの初期値
    let mut ans = INF;
    // 探索済みの頂点
    let mut used = vec![false; n];

    // DFS
    used[0] = true;
    dfs(0, n, &edge, &mut used, 0, &mut ans);
    used[0] = false;

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

// 1 -> Nの経路をDFSする。探索できた経路の総XORの最小値を求める。
fn dfs(
    cpos: usize,
    n: usize,
    edge: &Vec<Vec<(usize, usize)>>,
    used: &mut Vec<bool>,
    ret: usize, // 総XOR
    ans: &mut usize // 総XORの最小値
) {
    // Nに到達
    if cpos == n-1 {
        *ans = min(*ans, ret);
        return;
    }

    // 次の頂点を探す
    for &(npos, w) in &edge[cpos] {
        // 探索済みの頂点は見ない
        if used[npos] { continue;}

        // 次の頂点を探索。戻るときに探索済みのフラグを戻す
        used[npos] = true;
        dfs(npos, n, edge, used, ret ^ w, ans);
        used[npos] = false;
    }
}

E問題

問題

https://atcoder.jp/contests/abc396/tasks/abc396_e

解説

問題文では整数を扱う問題となっていますが、グラフ問題として解釈することができます。
( NN は頂点、MM は辺、XX は頂点の開始、YY は頂点の終点、ZZ は辺の値を表します。)
グラフ問題として考えた場合、頂点は辺を通ることで一部の頂点を辿ることができる連結グラフになります。したがって、連結グラフごとにDFSを行い、辿ることができた頂点に対して数列 AA の値を決定すれば良いです。
ここで、数列 AA の各値の決め方ですが、辺の値のXOR演算に依存しているので、以下の性質が使えます。

  • 0011 の二値情報のみで演算が出来る
  • 各ビットごとに演算が出来る

また、上記の性質を踏まえて各ビットごとに連結グラフを考えた場合、各頂点は 0011 のみ使用するため、二部グラフとみなすことができます。
したがって、各連結グラフについて各ビットごとのグラフに分解し、分解したグラフが全て二部グラフであれば、連結グラフに含まれる頂点の値を決定できます。
さらに、二部グラフでは 0011 を入れ替えることが可能なため、各二部グラフで 11 が少なくなるように選択していくことで、数列 AA の値を最小化することができます。

コード

abc396e.rs
use proconio::{input, marker::Usize1};
use itertools::Itertools;
const NOT_REACHED: usize = 255;

fn main() {
    // 入力
    input! {
        n: usize, m: usize,
        xyz: [(Usize1, Usize1, usize); m],
    }

    // 無向グラフの隣接リスト(向かいの頂点、辺の値)
    let mut edge = vec![Vec::new(); n];
    for (x, y, z) in xyz {
        edge[x].push((y, z));
        edge[y].push((x, z));
    }

    // 適切な二部グラフが作れるかどうかのフラグ
    let mut f = true;
    // 出力となる数列Aの値
    let mut ans = vec![0; n];

    // ビットごとに処理を行う
    for bi in 0..30 {
        // 頂点の探索済み有無(探索済みなら0か1、255は未探索)
        let mut used = vec![NOT_REACHED; n];

        // 開始頂点を探す
        for i in 0..n {
            // 既に見ている頂点はスキップ
            if used[i] != NOT_REACHED { continue; }
            // 一度の探索で見た0, 1のデータ
            let mut data01 = vec![Vec::new(); 2];

            // 初めの値は0としてセット
            used[i] = 0;
            data01[0].push(i);

            // DFSを実行
            f &= dfs(i, &edge, &mut used, bi, 0, &mut data01);
            // ビット毎の0と1の決定
            decide_01(&mut data01, &mut ans, bi);
        }
    }
    
    // 出力
    if !f { 
        println!("-1");
        return;
    }
    println!("{}", ans.iter().join(" "));
}

fn dfs(
    cpos: usize, 
    edge: &Vec<Vec<(usize, usize)>>,
    used: &mut Vec<usize>,
    bi: usize,
    cval: usize, // 現在のXOR値(0または1)
    data01: &mut Vec<Vec<usize>>, // 0と1のデータ
) -> bool {
    // 二部グラフ判定結果
    let mut f = true;

    // 次の頂点を探す
    for &(npos, val) in &edge[cpos] {
        // 次の頂点の値を決定
        let bit_val = if val & 1 << bi != 0 { 1 } else { 0 };
        let nval = cval ^ bit_val;

        // まだ見ていない頂点
        if used[npos] == NOT_REACHED {
            // 次の頂点の値を格納
            used[npos] = nval;
            data01[nval].push(npos);    
            // 次の頂点を探索
            f &= dfs(npos, edge, used, bi, nval, data01);
        }
        // 一度見た頂点
        else {
            // 二部グラフに矛盾が生じる。
            if used[npos] != nval {
                return false;
            }
            // 二部グラフに矛盾しない場合は、何もしない。
        }
    }
    f
}

fn decide_01(data01: &mut Vec<Vec<usize>>, ans: &mut Vec<usize>, k: usize) {
    // 1の個数が小さくなるようにスワップ
    let len0 = data01[0].len();
    let len1 = data01[1].len();
    if len0 < len1 {
        data01.swap(0, 1);
    }
    // 1が立っている箇所に1 << kを追加
    for &item in &data01[1] {
        ans[item] |= 1 << k;
    }
}

Discussion

ログインするとコメントできます