👨‍🍳

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

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

AtCoder Beginner Contest 416のA-E問題を緑コーダーが分かりやすく解説をまとめました。参考になりましたら幸いです。

ABC416-A

問題

https://atcoder.jp/contests/abc416/tasks/abc416_a

文字列 S の指定された区間 [L, R] が全て o かどうかを判定する問題です。

解説

指定された区間 [L, R] に含まれる文字を順に確認し、1つでも x が含まれていれば No 、全て o であれば、 Yes を出力します。

コード

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

fn main() {
    // 入力
    input! {
        _n: usize, // 文字列の長さ
        l: Usize1, r: Usize1, // 区間の開始と終了（0-index）
        s: Chars, // 文字列
    }

    // [l..r]にxが含まれていないかを判定
    let mut ans = true;
    for i in l..=r {
        if s[i] == 'x' { 
            ans = false;
        }
    }

    // 答えを出力
    println!("{}", if ans {"Yes"} else {"No"});
}

ABC416-B

問題

https://atcoder.jp/contests/abc416/tasks/abc416_b

文字列 S について、 o の文字数が最大となる配置を出力する問題です。

解説

問題文の条件を整理すると、以下のように解釈できます。

  • o と o の間には # を1つ以上配置する必要があります。
    → # と # で区切られた区間には、o は1つしか置けません。

この条件を満たすためには、文字列を # の固まりと . の固まりに分けて考え、以下の通りに出力します。

  • # の固まりは、そのまま出力します。
  • . の固まりは、先頭の1文字を o に変え、残りはそのまま . を出力します。

文字列を # の固まりと . の固まりに分けるには、ランレングス圧縮を用いるとよいです。

コード

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

fn main() {
    // 入力
    input! {
        s: Chars, // 文字列
    }

    // 文字列をランレングス圧縮する
    let runs = runlen(&s);

    // 文字の種類ごとに出力
    for (c, cnt) in runs {
        // #の場合は、そのまま文字数出力
        if c == '#' {
            for _ in 0..cnt { print!("#"); }
        }
        // .の場合は、先頭をoに変えて残り文字数はそのまま出力
        else {
            print!("o");
            for _ in 1..cnt { print!("."); }
        }
    }
    println!();
}

// ランレングス圧縮
pub fn runlen(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
}

ABC416-C

問題

https://atcoder.jp/contests/abc416/tasks/abc416_c

文字列を連結して作ることができる文字列について、辞書順で X 番目の文字列を出力する問題です。

解説

N 個の文字列から K 個を選んで連結する方法は、N^K 通り存在します。各連結文字列の長さは最大で50文字ですが、N と K の制約が小さいため、全ての組み合わせを総当たりで調べることが可能です。
具体的には以下1~4の手順で解きます。

  1. 再帰を用いて、K 個の文字列を選択する全ての組み合わせを生成します。
  2. 各組み合わせに対して選択した文字列を順番に連結して1つの文字列を作ります。
  3. 生成した全ての連結文字列を辞書順にソートします。
  4. ソート後、辞書順で X 番目の文字列を出力します。

コード

abc416c.rs
use itertools::Itertools;
use proconio::{input, marker::Usize1};

fn main() {
    input! {
        n: usize, // 文字列の個数
        k: usize, // つなげる個数
        x: Usize1, // 辞書順でx番目の文字列(0-index)
        s: [String; n], // 各文字列
    }

    // 連結文字列のリスト
    let mut concat_str_list = Vec::new();
    
    // 再帰を用いて全ての組み合わせを生成
    rec(n, 0, k, &mut concat_str_list, &mut vec![], &s);

    // 辞書順にソートしてx番目の文字列を出力
    concat_str_list.sort();
    println!("{}", concat_str_list[x].iter().join(""));
}

fn rec(
    n: usize, cnt: usize, // 文字列の個数, 現在選択中の個数
    k: usize, // つなげる個数
    concat_str_list: &mut Vec<Vec<char>>, // 生成した文字列リスト
    select: &mut Vec<usize>, // 選択順のリスト
    s: &Vec<String>, // 各文字列
) {
    // K個選択した場合
    if cnt == k {
        // 連結文字列を作成
        let mut ret = String::new();
        for &mut i in select {
            ret.push_str(&s[i]);
        }
        // リストに追加
        concat_str_list.push(ret.chars().collect());
    } else {
        // まだ選択が終わっていない場合
        for i in 0..n {
            // 文字列の番号を選択
            select.push(i);
            // 再帰で次の選択を調べる
            rec(n, cnt + 1, k, concat_str_list, select, s);
            // 選択した番号を戻す
            select.pop();
        }
    }
}

ABC416-D

問題

https://atcoder.jp/contests/abc416/tasks/abc416_d

(A_i + B_i) \mod M の総和を最小化する問題です。

解説

まず、(A_i + B_i) \mod M の総和は、以下のように分解できます。

\text{総和} = \text{Aの総和 + Bの総和} - C(※) \times M \\ \text{※} (A_i + B_i) \text{がM以上となる個数}

つまり、(A_i + B_i) が M 以上になるペアを多く作ることで、総和を最小化できます。

そのため、以下1~4の手順で解きます。

  1. 数列 A を昇順にソートします。
  2. 数列 B を降順にソートして、キューに追加します。
  3. A の値を小さい順に見ていき、その時点で残っている B の値の中で最も大きい値とペアを作り、キューの先頭を取り出します。
  4. そのペアの和が M 以上であれば、総和から M を引きます。

このようにして、M 以上のペアをできるだけ多く作ることで、総和を最小化します。

コード

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

fn main() {
    // テストケース数を入力
    input! {
        t: usize,
    }

    for _ in 0..t {
        solve();
    }
}

fn solve() {
    input! {
        n: usize, // 数列の長さ
        m: usize, // mod
        mut a: [usize; n], // 数列A
        mut b: [usize; n], // 数列B
    }

    // Aの総和 + Bの総和を求める
    let mut tot_ab: usize = a.iter().sum::<usize>() + b.iter().sum::<usize>();

    // Aは昇順にソート
    a.sort();

    // Bは降順にソートしてキューに積む
    b.sort_by(|i, j| j.cmp(i));
    let mut dq = VecDeque::from_iter(&b);

    // Aを全探索。a_i + b_i が M 以上なら、ペアにして取り除く
    for &aa in &a {
        if let Some(&bb) = dq.front() {
            if aa + bb >= m {
                tot_ab -= m;
                dq.pop_front();
            }
        }
    }

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

ABC416-E

問題

https://atcoder.jp/contests/abc416/tasks/abc416_e

N 個の街を頂点とし、道路や空路を辺とするグラフにおいて、最短経路を効率的に更新しながら、全頂点間の最短経路の総和を求める問題です。

解説

全頂点の最短経路の総和はワーシャルフロイドと呼ばれるアルゴリズムを用いることで解くことができます。ただし、道路や空路が増えていく問題のため、最短経路を更新する際の計算量に気を付ける必要があります。
具体的には、以下1~4の手順で解くことができます。

  1. 初期状態の最短経路の計算
    入力で与えられる道路と空路をもとに、初期状態の最短経路を計算します。この際、ワーシャルフロイド法を用いて全頂点間の最短経路を求めます。計算量は O(N^3) です。
    また、空路を扱いやすくするために、空港を束ねた「上空」という仮想的な頂点を追加し、空港と上空の間の最短経路を設定します。ただし、後述する空路の追加に備えて以下を行います。

    • 道路の最短経路は2倍の値をセットする。
    • 各空港と上空間の最短経路はそれぞれ T (往復の経路が2倍の値)をセットする。

  2. 道路の追加
    新たに追加される \text{(a,b)} の最短経路を更新します。
    この更新により影響を受ける以下2通りの経路順に通る場合について再計算を行います。

    • \text{a→b} の順に通る場合
    • \text{b→a} の順に通る場合

    計算量は O(N^2) です。

  3. 空路の追加
    追加する空港に対して、各空港の最短経路を更新すると、N 個の経路の追加と最短経路の再計算で O(N^3) の計算量となり、この更新では実行時間制限に間に合わなくなってしまいます。
    そこで、経路の更新を工夫して追加する空港と各空港から行き来できる上空、すなわち \text{(d,上空)} の最短経路を更新します。
    この更新も、以下2通りの経路順に通る場合について再計算を行います。

    • \text{d→上空} の順に通る場合
    • \text{上空→d} の順に通る場合

    計算量は O(N^2) です。

  4. 全頂点間の最短経路の総和の計算
    更新された最短経路をもとに、その時点で行き来可能な全頂点間の最短経路の総和を計算します。ただし、最短経路は2倍の値で管理しているため、総和を求めた後に 1/2 します。計算量は O(N^2) です。

初期状態の計算が O(N^3)、クエリ処理が O(N^2Q) であるため、全体の計算量は O(N^3 + N^2Q) となります。
問題文の制約では、N が最大500、 Q が最大1000のため、O(N^3 + N^2Q) = 3.75 \times 10^8 となり、実行時間制限以内に解くことが可能です。

コード

abc416e.rs
use std::cmp::min;
use proconio::{input, marker::Usize1};
const INF: usize = 1 << 60;

fn main() {
    // 入力
    input! {
        n: usize // 街の数
    }

    // 全頂点の最短経路(街の数+上空)を初期化
    let mut dist = vec![vec![INF; n + 1]; n + 1];
    for i in 0..=n {
        dist[i][i] = 0;
    }

    // 初期状態の道路の最短経路を反映
    input_init_road(&mut dist);

    // 初期状態の空路の最短経路を反映、空路の移動時間tを取得
    let mut t = 0;
    input_init_airroute(&mut dist, &mut t, n);

    // ワーシャルフロイド法による全頂点の最短経路を計算
    init_floyd_warshall(&mut dist, n + 1);
    
    // クエリ数
    input! {
        q: usize
    }

    // クエリ毎に処理
    for _ in 0..q {
        input! {
            nm: usize, 
        }
        // 道路の追加
        if nm == 1 {
            input! {
                a: Usize1, b: Usize1, c: usize, // 経路(a,b)、かかる時間c
            }
            // 追加する道路の最短経路を更新
            dist[a][b] = min(dist[a][b], 2 * c);
            dist[b][a] = min(dist[b][a], 2 * c);

            // 最短経路について、影響がある部分のみ更新
            update_dist(&mut dist, n + 1, a, b);
            update_dist(&mut dist, n + 1, b, a);
        }
        // 空港の追加
        else if nm == 2 {
            input! {
                d: Usize1, // 空港
            }
            // 追加する空路(空港->上空、上空->空港)の最短経路を更新
            dist[d][n] = min(dist[d][n], t);
            dist[n][d] = min(dist[n][d], t);
        
            // 最短経路について、影響がある部分のみ更新
            update_dist(&mut dist, n + 1, d, n);
            update_dist(&mut dist, n + 1, n, d);
        }
        // 全頂点の最短経路の総和を求める
        else {
            println!("{}", calc_total_dist(&dist, n));
        }
    }
}

fn input_init_road(dist: &mut Vec<Vec<usize>>) {
    input! {
        m: usize, // 道路の数
    }
    for _ in 0..m {
        input! {
            a: Usize1, b: Usize1, c: usize, // 経路(a,b)、かかる時間c
        }
        // a->b、b->aについて、かかる時間を2倍値でセットする
        dist[a][b] = min(dist[a][b], 2 * c);
        dist[b][a] = min(dist[b][a], 2 * c);
    }
}

fn input_init_airroute(dist: &mut Vec<Vec<usize>>, t: &mut usize, n: usize) {
    input! {
        k: usize, // 空港の個数
        tt: usize, // 所要時間
        d: [Usize1; k], // 空港の位置
    }

    for &dd in &d {
        // 空港->上空,上空->空港について、かかる時間をセットする
        dist[dd][n] = min(dist[dd][n], tt);
        dist[n][dd] = min(dist[n][dd], tt);
    }
    // 所要時間を返すためセット
    *t = tt;
}

fn init_floyd_warshall(dist: &mut Vec<Vec<usize>>, sz: usize) {
    for mid in 0..sz {
        for from in 0..sz {
            for to in 0..sz {
                dist[from][to] = min(dist[from][to], dist[from][mid] + dist[mid][to]);
            }
        }
    }
}

fn calc_total_dist(dist: &Vec<Vec<usize>>, n: usize) -> usize {
    let mut tot = 0;
    for i in 0..n {
        for j in 0..n {
            // 確定している経路のみ計上する
            if dist[i][j] < INF {
                tot += dist[i][j];
            }
        }
    }
    // 全距離を2倍にしているので、1/2する
    tot / 2
}

fn update_dist(dist: &mut Vec<Vec<usize>>, sz: usize, a: usize, b: usize) {
    for from in 0..sz {
        for to in 0..sz {
            dist[from][to] = min(dist[from][to], dist[from][a] + dist[a][b] + dist[b][to]);
            dist[from][to] = min(dist[from][to], dist[from][b] + dist[b][a] + dist[b][to]);
        }
    }
}

Discussion