Zenn
👨‍🍳

ABC401:Rustで解く！問題解説

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

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

A問題

問題

https://atcoder.jp/contests/abc401/tasks/abc401_a

解説

S の値が 200 以上 299 以下であれば、 Success を出力し、それ以外の場合は Failure を出力します。

コード

abc401a.rs
use proconio::input;

fn main() {
    // 入力
    input! {
        s: usize,
    }

    // 出力
    println!("{}", if s >= 200 && s <= 299 { "Success" } else { "Failure" });
}

B問題

問題

https://atcoder.jp/contests/abc401/tasks/abc401_b

解説

ログイン状態をフラグで管理しながら、与えられた操作を順にシミュレーションし、認証エラーが発生した( private の操作がログアウト状態で実行された)回数を求めます。

入力に基づき、以下の処理を行います：

  1. login または logout の入力を受け取った場合
    • login ならログイン状態を「有効」にします。
    • logout ならログイン状態を「無効」にします。
  2. private の入力を受け取った場合
    • ログイン状態が「無効」であれば認証エラーとし、エラー回数をカウントします。

コード

abc401b.rs
use proconio::input;

fn main() {
    // 入力
    input! {
        n: usize,
    }

    let mut is_login = false;
    let mut err_cnt = 0;

    for _ in 0..n {
        // 各操作を入力
        input! {
            s: String,
        }

        if s == "login" {
            is_login = true;
        } else if s == "logout" {
            is_login = false;
        } else if s == "private" && !is_login {
            err_cnt += 1;
        }
    }

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

C問題

問題

https://atcoder.jp/contests/abc401/tasks/abc401_c

解説

以下の条件を満たす数列 A = (A_0, A_1, \dots, A_N) の A_N を求めます。

  • 初めの K 項目までは 1 である。
  • K+1 項目以降は、直前の K 項の和である。

ただし、N, K の制約が 10^6 となっているため、直前の K 項の和をそのまま計算すると計算量が O(K) 、全体で O(NK) となり、実行時間制限以内に解くことができません。そのため、計算方法を工夫する必要があります。

直前の K 項の和を効率よく管理するために、累積和の考え方を利用します。具体的には、次のようにします。

  1. 初めの K 項はすべて 1 なので、その分の累積和はそのまま足して初期化します。
  2. K+1 項目以降は、直前の K 項の和を利用して計算します。これは、直前の合計から最も古い値( K 項前)を引いた上で、新たな値を加えることで効率よく求めることができます。

これにより、直前の K 項の和の計算量を O(1) に抑えることができます。

2.の計算について

まず、K 項目の総和を S_k = a_1 + a_2 + \dots + a_k としたとき、
K+1 項目の値、K+1 項目の総和は以下となります。

a_{k+1} = S_k
\begin{aligned} S_{k+1} &= a_2 + \dots + a_k + a_{k+1} \\ &= a_2 + \dots + a_k + S_k \\ \end{aligned}

ここで、K+1 項目の総和と K 項目の総和の差分を考えると、

\begin{aligned} S_{k+1} - S_k　&= (a_2 + ... + a_k + a_{k+1}) - (a_1 + a_2 + ... + a_{k}) \\ &= a_{k+1} - a_1 \\ &= S_k - a_1 \\ \end{aligned}

となります。
よって、 K+1 項目の総和は、K 項目の総和 + S_k - a_1 で求めることができます。

また、答えは 10^9 で割った余りを求める必要があるため、計算中に余りを取ることでオーバーフローを防ぎます。

コード

abc401c.rs
use proconio::input;

const MOD: usize = 1_000_000_000;

fn main() {
    // 入力
    input! {
        n: usize,
        k: usize,
    }

    let mut dp = vec![0; n + 1];
    let mut tot = 0;

    for i in 0..=n {
        if i < k {
            // 初めのK項は1
            dp[i] = 1;
            tot += dp[i];
            tot %= MOD;
        } else {
            // K+1項目以降
            dp[i] = tot;
            tot += dp[i] + MOD;
            tot -= dp[i - k];
            tot %= MOD;
        }
    }

    println!("{}", dp[n]);
}

D問題

問題

https://atcoder.jp/contests/abc401/tasks/abc401_d

解説

この問題では、長さ N の文字列 S に含まれる ? を o または . に置き換えることで、以下の条件を満たすようにします。

oの個数がちょうどK個
oが連続しない

この問題を解くために、以下の手順で処理を行います。

  1. 前処理
    文字列中の o の両隣にある ? を . に置き換えて、o が連続しないようにします。

  2. ランレングス圧縮
    文字列を連続する文字ごとにまとめ、(文字, 長さ) のペアのリストを作成します。これにより、? の連続部分や o の個数を効率的に処理できます。

  3. 現在の o の個数と最大可能な o の個数を計算
    ランレングス圧縮した結果を基に、確定している o の個数と、? を最大限 o に変えた場合の最大 o の個数を計算します。

  4. 場合分け
    以下の3つのケースに分けて処理します。

(1)確定している o の個数が K と一致する場合
→ ? をすべて . に置き換えて出力します。

(2)確定している o の個数と最大可能な o の個数の合計が K を超える場合
→置き換えなくてよい o が存在するので、o の決め方が一意に定まりません。そのため、入力文字列の ? は変更せずにそのまま出力します。

(3)確定している o の個数と最大可能な o の個数の合計がちょうど K の場合
→ ? の連続部分毎の個数によって決まります。

(3)-1. ? が連続する個数が奇数の場合
→ 最初を o として、交互に o と . を割り当てると適切に置き換えができるので、? を o または . に置き換えて出力します。

(3)-2. ? が連続する個数が偶数の場合
→o から始める場合と . が始める場合の2パターン考えられるので、o の決め方が一意に定まりません。そのため、入力文字列の ? は変更せずにそのまま出力します。

これらの手順を整理することで、問題を効率的に解くことができます。

コード

abc401d.rs
use proconio::{input, marker::Chars};

fn main() {
    input! {
        n: usize,
        k: usize,
        s: Chars,
    }

    let mut ss = s.clone();

    // 前処理: 'o' の両隣の '?' を '.' に変える
    pre_change(&mut ss, n);

    // ランレングス圧縮: 文字列を圧縮して (文字, 長さ) のペアのリストを作成
    let rangs = ranglen(&ss);

    // 確定している 'o' の個数を数える
    let confirmed_o_cnt = get_confirm_o_cnt(&rangs);

    // '?' を最大限 'o' に変えた場合の最大 'o' の個数を計算
    let possible_o_max_cnt = get_possible_o_max_cnt(&rangs);

    // 場合分けして出力
    if confirmed_o_cnt == k {
        // 確定している 'o' の数がちょうど k の場合
        print_alternating_dot(&rangs);
    } else if confirmed_o_cnt + possible_o_max_cnt > k {
        // 確定している 'o' と '?' を最大限 'o' に変えると、 k を超える場合
        print_compressed_string(&rangs);
    } else if confirmed_o_cnt + possible_o_max_cnt == k {
        // 確定している 'o' と '?' を最大限 'o' に変えてちょうど k の場合
        print_alternating_o_and_dot(&rangs);
    }
}

// 前処理: 'o' の両隣の '?' を '.' に変える
fn pre_change(string: &mut Vec<char>, length: usize) {
    for i in 0..length {
        if string[i] == 'o' {
            if i > 0 && string[i - 1] == '?' {
                string[i - 1] = '.';
            }
            if i < length - 1 && string[i + 1] == '?' {
                string[i + 1] = '.';
            }
        }
    }
}

// ランレングス圧縮: 文字列を (文字, 長さ) のペアのリストに変換
fn ranglen(arr: &Vec<char>) -> Vec<(char, usize)> {
    let len = arr.len();
    let mut arr_ret = Vec::new();
    let mut tail = 0;
    let mut top = 0;
    while tail < len {
        let now_c = arr[tail];
        while top < len && now_c == arr[top] {
            top += 1;
        }
        arr_ret.push((now_c, top - tail));
        tail = top;
    }
    arr_ret
}

// 確定している 'o' の個数を数える
fn get_confirm_o_cnt(rangs: &Vec<(char, usize)>) -> usize {
    let mut ret = 0;
    for &(c, l) in rangs {
        if c == 'o' { ret += l; }
    }
    ret
}

// '?' を最大限 'o' に変えた場合の最大 'o' の個数を計算
fn get_possible_o_max_cnt(rangs: &Vec<(char, usize)>) -> usize {
    let mut ret = 0;
    for &(c, l) in rangs {
        // '?' を交互に 'o' に変える
        if c == '?' { ret += (l + 1) / 2; }
    }
    ret
}

// 圧縮された文字列をそのまま出力
fn print_compressed_string(rangs: &Vec<(char, usize)>) {
    for &(char, length) in rangs {
        for _ in 0..length {
            print!("{}", char);
        }
    }
    println!();
}

// '?' を'.'に置き換えて出力
fn print_alternating_dot(rangs: &Vec<(char, usize)>) {
    for &(char, length) in rangs {
        if char == '?' {
            for _ in 0..length {
                print!("{}", '.');
            }
        } else {
            for _ in 0..length {
                print!("{}", char);
            }
        }
    }
    println!();
}

// '?' を交互に 'o' と '.' に変えて出力
fn print_alternating_o_and_dot(rangs: &Vec<(char, usize)>) {
    for &(char, length) in rangs {
        if char == '?' && length % 2 == 1 {
            for i in 0..length {
                print!("{}", if i % 2 == 0 { 'o' } else { '.' });
            }
        } else {
            for _ in 0..length {
                print!("{}", char);
            }
        }
    }
    println!();
}

E問題

問題

https://atcoder.jp/contests/abc401/tasks/abc401_e

解説

この問題は、頂点1から頂点 K (K=1, \dots, N) まで全て連結した状態を満たしたまま、頂点 K+1 以降の不要な頂点をいくつ取り除けるかを求める問題です。これはUnion-Findというデータ構造を活用して解くことができます。

具体的には、以下の2つのUnion-Find構造を用意します。

  1. 頂点1から頂点 K までを連結するためのUnion-Find。以下uf_smallとします。
  2. 頂点1から順番に見ていき、現時点で連結可能な頂点を管理するUnion-Find。以下uf_largeとします。

アルゴリズムの流れは以下の通りです：

  1. 各頂点について、その頂点から辿れる辺を順次処理します。
    • uf_largeでは、現在の頂点番号より大きい番号の頂点とマージします。
    • uf_smallでは、現在の頂点番号以下の頂点とマージします。
  2. uf_smallが、頂点1から現在の頂点まで全て連結されていない場合、-1を出力します。
    そうでない場合、uf_largeの連結サイズからuf_smallの連結サイズを引いた値が、不要な頂点の数となるので、その値を出力します。

コード

abc401e.rs
use proconio::{input, marker::Usize1};
use std::mem::swap;

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

    // 無向グラフの辺
    let mut edges = vec![Vec::new(); n];
    for _ in 0..m {
        input! {
            u: Usize1, v: Usize1,
        }
        edges[u].push(v);
        edges[v].push(u);
    }

    // Union-Findを2つ用意
    let mut uf_small = UnionFind::new(n);
    let mut uf_large = UnionFind::new(n);

    for cpos in 0..n {
        for &npos in &edges[cpos] {
            if cpos < npos {
                uf_large.merge(cpos, npos);
            } else {
                uf_small.merge(cpos, npos);
            }
        }
        // uf_smallは1~cposまで全て連結でないといけない
        if uf_small.size(cpos) != cpos + 1 {
            println!("-1");
        } 
        // uf_largeの連結サイズからuf_smallの連結サイズを引いたものが不要な頂点
        else {    
            println!("{}", uf_large.size(cpos) - uf_small.size(cpos));
        }
    }
}

struct UnionFind {
    size: Vec<usize>,        // 各頂点のグラフサイズ
    parent: Vec<isize>,      // 親情報
}
impl UnionFind {
    fn new(n : usize) -> Self {
        UnionFind { size : vec![1; n], parent : vec![-1; n]}
    }
    // 代表値を求める。(経路圧縮も行う)
    #[allow(unused)]
    fn find(&mut self, u: usize) -> usize {
        let mut mut_u = u;
        if self.parent[mut_u] == -1 { return mut_u; }

        let p = self.parent[mut_u] as usize;

        let par = self.find(p);
        self.parent[mut_u] = par as isize;
        return self.parent[mut_u] as usize;    
    }
    // 併合
    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);
        }
        // yをxの子とする
        self.parent[y] = x as isize;
        self.size[x] += self.size[y];
        true
    }
    // 同じ集合かどうかを判定
    #[allow(unused)]
    fn same(&mut self, u: usize, v: usize) -> bool {
        self.find(u) == self.find(v)
    }
    // 集合のサイズを求める
    #[allow(unused)]
    fn size(&mut self, u: usize) -> usize {
        let x = self.find(u);
        self.size[x]
    }
}

Discussion