👨‍🍳

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

に公開

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

ABC429-A

問題

https://atcoder.jp/contests/abc429/tasks/abc429_a

N 回のリクエストに対する応答結果を出力する問題です。

解説

リクエストを順番に処理し、リクエスト回数が指定された上限 M 回以下の場合は OK を出力し、M 回を超える場合は Too Many Requests を出力します。

コード

abc429a.rs
use proconio::input;

fn main() {
    // 入力
    input! {
        n: usize, // リクエスト回数
        m: usize, // レスポンス回数の上限
    }

    // リクエストを順番に処理
    for i in 0..n {
        // M回以内なら正常レスポンス
        if i < m {
            println!("OK");
        }
        // M回を超えたらエラーレスポンス
        else {
            println!("Too Many Requests");
        }
    }
}

ABC429-B

問題

https://atcoder.jp/contests/abc429/tasks/abc429_b

数列 A について、1つを除いて合計した値が M と一致するかを判定する問題です。

解説

まず、数列 A の総和を計算します。これを S とします。
次に、数列 A の各要素を1つずつ取り除いた場合の合計値を計算します。この合計値は S - A[i] となります。
この合計値が M と一致するかを判定します。一致する場合は Yes を出力し、そうでない場合は No を出力します。

コード

abc429b.rs
use proconio::input;

fn main() {
    input! {
        n: usize, // 数列の長さ
        m: usize, // 合計値
        a: [usize; n], // 数列の値
    }

    // 数列Aの総和を求める
    let tot_a: usize = a.iter().sum();

    // 総和から1つ除いた合計値がMと一致するかをチェック
    for &val in &a {
        if tot_a - val == m {
            println!("Yes");
            return;
        }
    }
    println!("No");
}

ABC429-C

問題

https://atcoder.jp/contests/abc429/tasks/abc429_c

数列 A について任意の値を3つ選ぶとき、2つの値が同じとなる組み合わせは何通りかを答える問題です。

解説

数列 A の中で、各値が何回出現するかを管理します。これには連想配列(HashMap)を用いると便利です。
ある値 XK 個存在する場合、次のように組み合わせを計算します。

  • ある値 X から2つを選ぶ組み合わせは、\binom{K}{2} = \frac{K (K-1)}{2} 通りです。
  • 残りの値(X 以外)から1つを選ぶ組み合わせは、N - K 通りです。
    したがって、X を2つ選び、残りの値から1つ選ぶ組み合わせは、\binom{K}{2} \times (N - K) 通りとなります。

上記の計算をすべての値について行い、総和を求めます。

コード

abc429c.rs
use std::collections::HashMap;
use proconio::input;

fn main() {
    input! {
        n: usize, // 数列の長さ
        a: [usize; n], // 数列の値
    }

    // 数列を値毎にまとめる
    let mp = collect_val(n, a);

    // 組み合わせの総数
    let mut ans = 0;

    // 2つ取り出す値を順番に調べる
    for (_val, cnt) in mp {
        // 同じ値から2つの選び方
        let choose2 = cnt * (cnt - 1) / 2;
        // 残りの値の選び方
        let other = n - cnt;
        // 値を2つ * 残り1つの選び方
        ans += choose2 * other;
    }
    // 組み合わせの総数を出力
    println!("{}", ans);
}

// 数列を値毎にまとめる関数
fn collect_val(_n: usize, a: Vec<usize>) -> HashMap<usize, usize> {
    let mut mp = HashMap::new();
    for &aa in &a {
        *mp.entry(aa).or_insert(0) += 1;
    }
    mp
}

ABC429-D

問題

https://atcoder.jp/contests/abc429/tasks/abc429_d

ある地点から円周上を時計回りに移動して C 人以上の人を見つける操作を、 M 個の地点で行ったときの総和を求める問題です。

解説

この問題では、以下1.~3.について考えることで解くことができます。

  1. 人がいる地点に注目する
    M 個の地点は非常に大きい (10^{12}) ため、人がいる地点だけに注目し、その地点の座標と人数を整理します。

  2. 円周上の性質を考慮する
    円周上の地点がループする性質を考慮する必要があります。
    そのため、2周分の情報を用意しておき、1周分を調べるようにします。

  3. 差分更新を利用する
    開始地点を決めて、そこから C 人以上見つけるまでの範囲を計算します。
    その後、次の地点に移動する際には、前の地点の人を除き、新しい地点の人を追加する形で差分更新を行います。

これを、各地点について、「範囲内の移動距離(現在の地点と次の地点との間の距離)」と「範囲内の人数」を掛け合わせることで、求める総和が得られます。

コード

abc429d.rs
use std::collections::{BTreeMap, BTreeSet};
use proconio::input;

#[derive(PartialEq, Eq, PartialOrd, Ord, Hash, Clone)]
struct PersonInfo {
    pos: usize,
    cnt: usize,
}

fn main() {
    input! {
        n: usize, // 人の人数
        m: usize, // 地点の数
        c: usize, // 移動条件(c以上になったら止まる)
        mut a: [usize; n], // 人の位置(0-index)   
    }
    // 人の位置を昇順にソート
    a.sort();

    // 地点ごとの人数情報を取得 (2周分)
    let person_double = get_person_double(n, m, &a);
    // 地点リストを取得 (2周分)
    let pos_double = get_pos_double(n, m, &a);
    // 地点リストの1周分のサイズ
    let sz = pos_double.len() / 2;

    // 総和を計算する変数
    let mut ans = 0;
    // 現在の人数
    let mut cur_person = 0;

    // 調べる範囲の[tail..head]について、開始区間を1周分調べる
    let mut head = 1;
    for tail in 1..=sz {
        // $C$ 人を超えるまで範囲を広げる
        while cur_person < c {
            cur_person += person_double[head].cnt;
            head += 1;
        }
        // 範囲内の移動距離 × 範囲内の人数を加算
        ans += (pos_double[tail] - pos_double[tail-1]) * cur_person;
        // 範囲の先頭地点の人を除外
        cur_person -= person_double[tail].cnt;
    }
    // 総和を出力
    println!("{}", ans);
}

// 地点ごとの人数情報を取得 (2周分)
fn get_person_double(_n: usize, m: usize, a: &Vec<usize>) -> Vec<PersonInfo> {
    let mut mp2 = BTreeMap::new();
    for &aa in a {
        *mp2.entry(aa).or_insert(0) += 1;
        *mp2.entry(aa + m).or_insert(0) += 1;
    }

    let mut posinfo = Vec::new();
    for (pos, cnt) in mp2 {
        posinfo.push( PersonInfo{pos, cnt});
    }
    posinfo
}

// 地点リストを取得 (2周分)
fn get_pos_double(_n: usize, m: usize, a: &Vec<usize>) -> Vec<usize> {
    let mut st = BTreeSet::new();
    for &aa in a {
        st.insert(aa);
        st.insert(aa + m);
    }

    let mut pos_list = Vec::new();
    for pos in st {
        pos_list.push(pos);
    }
    pos_list
}

ABC429-E

問題

https://atcoder.jp/contests/abc429/tasks/abc429_e

安全な頂点 A -> 危険な頂点 -> 安全な頂点 B の最短距離を求める問題です。

解説

幅優先探索(BFS)を用いて解くことができます。
具体的な考え方は以下1.~3.の通りです。

  1. 安全な頂点を始点として探索
    安全な頂点をすべて始点として、危険な頂点までの最短距離をBFSで求めます。

  2. 危険な頂点に対する最短距離の記録
    各危険な頂点について、最短距離が1番目に近い安全な頂点と、2番目に近い安全な頂点を記録します。

  3. 距離の和を計算
    記録した1番目と2番目の距離を足し合わせることで、問題の条件を満たす最短距離を求めます。

コード

abc429e.rs
use std::collections::VecDeque;
use proconio::{input, marker::{Chars, Usize1}};
const INF: usize = 1 << 60;

#[derive(PartialEq, Eq, PartialOrd, Ord, Hash, Clone)]
struct MinDist {
    pos: usize,
    dist: usize,
}

fn main() {
    input! {
        n: usize, // 頂点数
        m: usize, // 辺数
        uv: [(Usize1, Usize1); m], // 辺 (u, v) (0-indexed)
        s: Chars, // 頂点の安全 ('S')、危険 ('D') 情報
    }
    // 無向グラフを作成
    let graph = make_graph(n, m, uv);
    // BFSで最短距離を計算
    let ans = bfs(n, graph, s);
    // 結果を出力
    for dist in ans {
        println!("{}", dist);
    }
}

fn make_graph(n: usize, _m: usize, uv: Vec<(usize, usize)>) -> Vec<Vec<usize>> {
    let mut ret = vec![Vec::new(); n];
    for (u, v) in uv {
        ret[u].push(v);
        ret[v].push(u);
    }
    ret
}

fn bfs(n: usize, graph: Vec<Vec<usize>>, s: Vec<char>) -> Vec<usize> {
    // 安全な頂点リストを取得
    let safe_pos = get_safe_pos(n, &s);

    // 1番目、2番目に最短な頂点と距離
    let mut min_first = vec![MinDist {pos: INF, dist: INF}; n];
    let mut min_second = vec![MinDist {pos: INF, dist: INF}; n];
    // 探索する頂点のキュー
    let mut dq = VecDeque::new();

    // キューと最短距離情報を初期化
    for &pos in &safe_pos {
        min_first[pos].pos = pos;
        min_first[pos].dist = 0;
        dq.push_back((pos, pos, 0)); // 開始位置、現在位置、最短距離をキューにセット
    }

    // BFSを実行
    while let Some((spos, cpos, dist)) = dq.pop_front() {
        for &npos in &graph[cpos] {
            // 1番目に最短な距離
            if min_first[npos].dist == INF {
                min_first[npos].pos = spos;
                min_first[npos].dist = dist + 1;
                dq.push_back((spos, npos, dist + 1));
            }
            // 2番目に最短な距離
            else if min_second[npos].dist == INF && min_first[npos].pos != spos {
                min_second[npos].pos = spos;
                min_second[npos].dist = dist + 1;
                dq.push_back((spos, npos, dist + 1));
            }
        }
    }

    // 危険な頂点について、安全な頂点を2つ選んで距離の和を求める
    let mut ret = Vec::new();
    for pos in 0..n {
        if s[pos] == 'D' {
            let dist = min_first[pos].dist + min_second[pos].dist;
            ret.push(dist);
        }
    }
    ret
}

fn get_safe_pos(n: usize, s: &Vec<char>) -> Vec<usize> {
    let mut ret = Vec::new();
    for i in 0..n {
        if s[i] == 'S' {
            ret.push(i);
        }
    }
    ret
}

Discussion