緑コーダーがRustで解説してみた(ABC420 A～E)

AtCoder Beginner Contest 420のA～E問題を緑コーダーが自分なりの解説と解答コードをまとめました。
参考になりましたら幸いです。

ABC420-A

問題

https://atcoder.jp/contests/ABC420/tasks/ABC420_a

$X$ 月の $Y$ カ月後の月を求める問題です。

解説

現在の月 $X$ に $Y$ カ月を足して求めます。ただし、$X + Y$ が12を超える場合は、年をまたぐため、12を引いて調整します。

コード

abc420a.rs

use proconio::input;

fn main() {
    // 入力
    input! {
        x: usize, // 現在の月
        y: usize, // 加算する月数
    }

    let mut nx = x + y;

    // 年をまたいだ場合は12を引く
    if nx > 12 {
        nx -= 12;
    }

    // 答えを出力
    println!("{}", nx);
}

ABC420-B

問題

https://atcoder.jp/contests/ABC420/tasks/ABC420_b

少数派に投票をすると得点が入るという問題です。

解説

各人の投票結果は縦に並んでいるので、1列ずつ調べていきます。各ラウンドごとに全員の投票を確認して、 0 と 1 のどちらが少数派かを判定し、少数派に投票した人に1点を加算します。
全ラウンド終了後、得点が最も高い人を出力します。

コード

abc420b.rs

use proconio::{input, marker::Chars};

fn main() {
    // 入力
    input! {
        n: usize, // 人
        m: usize, // ラウンド数
        s: [Chars; n], // 投票情報
    }

    // 各人のスコア
    let mut score = vec![0; n];

    // ラウンド毎に見る
    for round in 0..m {
        // 0と1の個数を集計して、少数派の値を求める
        let minority = get_minority_num(n, &s, round);
        // 少数派の人に得点を加算
        update_score(n, &s, round, &mut score, minority);
    } 

    // スコアが最も高い人を全員出力
    let max_score = *score.iter().max().unwrap();
    for i in 0..n {
        if score[i] == max_score {
            print!("{} ", i + 1); // 1-indexで出力
        }
    }
    println!();
}

// 少数派の値を判定する関数
fn get_minority_num(n: usize, s: &Vec<Vec<char>>, round: usize) -> usize {
    let mut cnt0 = 0;
    let mut cnt1 = 0;
    for i in 0..n {
        if s[i][round] == '0' { 
            cnt0 += 1; 
        } else { 
            cnt1 += 1; 
        }
    }

    if cnt0 < cnt1 { 0 } else { 1 }
}

// 少数派の人に得点を加算する関数
fn update_score(n: usize, s: &Vec<Vec<char>>, round: usize, score: &mut Vec<usize>, minority: usize) {
    for i in 0..n {
        if minority == 0 && s[i][round] == '0' { 
            score[i] += 1; 
        }
        if minority == 1 && s[i][round] == '1' { 
            score[i] += 1; 
        }
    }
}

ABC420-C

問題

https://atcoder.jp/contests/ABC420/tasks/ABC420_c

数列A, Bの要素ごとの最小値の総和を求める問題です。

解説

各クエリで数列AまたはBの値が変更されるたびに、対応するインデックスの最小値を更新し、最小値の総和を効率的に計算する必要があります。最小値の総和は、以下1~3のようにして変更されたインデックスの差分のみを更新します。

1. 更新前のインデックスの最小値を引きます。
2. 更新箇所のインデックスの新しい最小値を求めます。
3. 更新後のインデックスの最小値を加えます。

これにより、クエリごとの計算量が $O(1)$ となり、全体の計算量は$O(N + Q)$ となります。

コード

abc420c.rs

use proconio::{input, marker::Usize1};
use std::cmp::min;

fn main() {
    // 入力
    input! {
        n: usize, // 長さN
        q: usize, // クエリ数
        mut a: [i64; n], // 数列A
        mut b: [i64; n], // 数列B
    }

    // 数列a,Bのminリスト
    let mut ab_min = vec![0; n];
    // 数列A,Bのminの総和
    let mut sum_min = 0;

    // minリスト,minの総和を初期化
    for i in 0..n {
        ab_min[i] = min(a[i], b[i]);
        sum_min += ab_min[i];
    }

    // クエリ処理
    for _ in 0..q {
        input! {
            c: char, // 変更する数列(A,B)
            x: Usize1, // 変更するindex
            y: i64, // 変更する値
        }

        // 変更前のminを引く
        sum_min -= ab_min[x];

        // 値の更新
        if c == 'A' {
            a[x] = y;
        }
        else {
            b[x] = y;
        }
        ab_min[x] = min(a[x], b[x]);
        
        // 変更後のminを足す
        sum_min += ab_min[x];
        
        // 更新後の値を出力
        println!("{}", sum_min);
    }
}

ABC420-D

問題

https://atcoder.jp/contests/ABC420/tasks/ABC420_d

スイッチマスの通過でドアの開閉が行われるグリッド上で、スタートマスからゴールマスへの最短距離を求める問題です。

解説

幅優先探索(BFS)を用いることで解くことができます。
スイッチマス ? を通過すると、現在の状態（スイッチを押しているか押していないか）が反転します。スイッチを押していない状態(はじめの状態)を「表面」、スイッチを押した状態を「裏面」として扱い、スタートマスからゴールマスまでの最短距離を求めます。通常のマス . では、状態を維持したまま移動し、スイッチマス ? では、状態を反転させて移動します。
ゴールマスに到達した際、表面と裏面のうち、より短い距離を答えとします。なお、ゴールマスに到達できない場合は -1 を出力します。

コード

abc420d.rs

use proconio::{input, marker::Chars};
use std::collections::VecDeque;
use std::cmp::min;

const INF: usize = 1 << 60;
const D4: [Position; 4] = [Position{x:!0, y:0}, Position{x:0, y:!0}, Position{x:1, y:0}, Position{x:0, y:1}]; // 上左下右

#[derive(Clone)]
struct Position {
    x: usize,
    y: usize,
}

fn main() {
    input! {
        h: usize, // グリッドの行
        w: usize, // グリッドの列
        s: [Chars; h], // グリッド
    }

    // スイッチoff/onのグリッド(2面)を初期化
    let mut grid = vec![vec![vec!['.'; 2]; w]; h];
    init_grid(h, w, &s, &mut grid);

    // スタートとゴールの位置を取得
    let spos = get_pos(h, w, &s, 'S');
    let gpos = get_pos(h, w, &s, 'G');

    // BFS
    let dist_min = bfs(h, w, &grid, spos, gpos);

    // ゴールに到達できない場合は-1
    if dist_min == INF {
        println!("-1");
    } else {
        // ゴールに到達できた場合はその値を出力
        println!("{}", dist_min);
    }
}

fn init_grid(h: usize, w: usize, s: &Vec<Vec<char>>, grid: &mut Vec<Vec<Vec<char>>>) {
    for i in 0..h {
        for j in 0..w {
            // #や?は、on/off共通でその文字を入れる
            if s[i][j] == '#' || s[i][j] == '?' {
                grid[i][j][0] = s[i][j];
                grid[i][j][1] = s[i][j];
            }
            // xは表面なら、#
            if s[i][j] == 'x' {
                grid[i][j][0] = '#';
            }
            // oは裏面なら、#
            if s[i][j] == 'o' {
                grid[i][j][1] = '#';
            }
        }
    }
}

fn get_pos(h: usize, w: usize, s: &Vec<Vec<char>>, c: char) -> Position {
    for i in 0..h {
        for j in 0..w {
            if s[i][j] == c {
                return Position { x: i, y: j };
            }
        }
    }
    Position { x: 0, y: 0 }
}

fn bfs(
    h: usize,
    w: usize,
    grid: &Vec<Vec<Vec<char>>>,
    spos: Position,
    gpos: Position,
) -> usize {
    // 最短距離の初期化
    let mut dist = vec![vec![vec![INF; 2]; w]; h];
    dist[spos.x][spos.y][0] = 0;

    // キューの初期化
    let mut dq = VecDeque::new();
    dq.push_back((spos, 0));

    // BFSを実施
    while let Some((cur_pos, cur_state)) = dq.pop_front() {
        // 4方向移動を試す
        for di in D4 {
            let mut next_pos = Position { x: 0, y: 0 };
            next_pos.x = cur_pos.x.wrapping_add(di.x);
            next_pos.y = cur_pos.y.wrapping_add(di.y);

            // グリッドの範囲外へは移動不可
            if next_pos.x >= h || next_pos.y >= w {
                continue;
            }
            // 壁へ移動は不可
            if grid[next_pos.x][next_pos.y][cur_state] == '#' {
                continue;
            }

            // .のマスへ移動する場合
            if grid[next_pos.x][next_pos.y][cur_state] == '.' {
                if dist[cur_pos.x][cur_pos.y][cur_state] + 1 < dist[next_pos.x][next_pos.y][cur_state] {
                    dist[next_pos.x][next_pos.y][cur_state] = dist[cur_pos.x][cur_pos.y][cur_state] + 1;
                    dq.push_back((next_pos.clone(), cur_state));
                }
            }

            // ?のマスへ移動する場合
            if grid[next_pos.x][next_pos.y][cur_state] == '?' {
                let next_state = cur_state ^ 1;
                if dist[cur_pos.x][cur_pos.y][cur_state] + 1 < dist[next_pos.x][next_pos.y][next_state] {
                    dist[next_pos.x][next_pos.y][next_state] = dist[cur_pos.x][cur_pos.y][cur_state] + 1;
                    dq.push_back((next_pos.clone(), next_state));
                }
            }
        }
    }

    // 表面, 裏面のうち、より小さい最短距離を返す
    min(dist[gpos.x][gpos.y][0], dist[gpos.x][gpos.y][1])
}

ABC420-E

問題

https://atcoder.jp/contests/ABC420/tasks/ABC420_e

N頂点の無向グラフにおいて、黒色の頂点に到達できるかどうかを判定する問題です。

解説

Union-Find（素集合データ構造）とHashSet(集合)を用いて、以下の3種類のクエリを効率的に処理します。

1. 頂点 $u$ と $v$ を結ぶ無向辺を追加する。
   頂点 $u$ と $v$ を連結します。この際、連結成分の親ノードが持つ黒色頂点の集合もマージします。

2. 頂点 $v$ の色を白から黒、または黒から白に変更する。
   頂点 $v$ の色を反転させます。この際、連結成分の親ノードが持つ黒色頂点の集合も更新します。

3. 頂点 $v$ が所属する連結成分に黒色の頂点が含まれるかを判定する。
   頂点 $v$ の連結成分に黒色の頂点が含まれるかを判定します。これは、連結成分の親ノードが持つ黒色頂点の集合に要素が1つ以上あるかを確認することで判断できます。

コード

abc420e.rs

use proconio::{input, marker::Usize1};
use std::collections::HashSet;
use std::mem::swap;

fn main() {
    // 入力
    input! {
        n: usize, // 頂点
        q: usize, // クエリ数
    }

    // Union-Findの初期化
    let mut uf = UnionFind::new(n);

    // 各グループの親が持っている黒頂点の集合を管理
    let mut black = vec![HashSet::new(); n];

    // クエリ処理
    for _ in 0..q {
        input! {
            nm: usize, // クエリ番号
        }
        if nm == 3 {
            input! {
                v: Usize1, // 頂点v(0-index)
            }
            // 黒頂点が1つ以上あるかどうかを判定する
            let leader = uf.find(v);
            println!("{}", if black[leader].len() >= 1 {"Yes"} else {"No"});		
        }
        else if nm == 2 {
            input! {
                v: Usize1, // 頂点v(0-index)
            }
            // 頂点vの白と黒を入れ替え。各グループの親に反映する
            let leader = uf.find(v);	
            if black[leader].contains(&v) {
                black[leader].remove(&v);
            }
            else {
                black[leader].insert(v);
            }
        }
        else {
            input! {
                mut u: Usize1, // 頂点u(0-index)
                mut v: Usize1, // 頂点v(0-index)
            }
            let leader_u = uf.find(u);
            let leader_v = uf.find(v);

            // 頂点uとvをマージ
            if uf.merge(u, v) {
                // 黒頂点集合のマージ
                if uf.find(u) == leader_u {
                    // 親がuの場合は、vをマージ
                    let tmp = std::mem::take(&mut black[leader_v]);
                    black[leader_u].extend(tmp);
                }
                else {
                    // 親がvの場合は、uをマージ
                    let tmp = std::mem::take(&mut black[leader_u]);
                    black[leader_v].extend(tmp);
                }
            }
        }
    }
}

// Union-Find構造体
struct UnionFind {
    size: Vec<usize>,        // 各頂点の集合サイズ
    parent: Vec<isize>,      // 親情報
}

impl UnionFind {
    // 初期化
    fn new(n: usize) -> Self {
        UnionFind { size: vec![1; n], parent: vec![-1; n] }
    }

    // 代表元を求める
    fn find(&mut self, u: usize) -> usize {
        if self.parent[u] == -1 {
            return u;
        }
        let p = self.parent[u] as usize;
        let root = self.find(p);
        // 経路圧縮
        self.parent[u] = root as isize;
        root
    }

    // 2つの集合を併合
    fn merge(&mut self, u: usize, v: usize) -> bool {
        let mut x = self.find(u);
        let mut y = self.find(v);
        if x == y {
            // 既に同じ集合
            return false; 
        }
        if self.size[x] < self.size[y] {
            swap(&mut x, &mut y);
        }
        self.parent[y] = x as isize;
        self.size[x] += self.size[y];
        true
    }

    // 集合のサイズを取得
    #[allow(unused)]
    fn size(&mut self, u: usize) -> usize {
        let root = self.find(u);
        self.size[root]
    }
}