👨‍🍳

ABC409: Rustで解く!問題解説

に公開

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

A 問題

問題

https://atcoder.jp/contests/abc409/tasks/abc409_a

解説

2人の欲しいものリスト(文字列)を比較し、同じ位置で両方が o となっている列が1つでも存在するかを判定する問題です。
具体的には、文字列 TA の各文字を順に比較し、どちらも o である箇所が1つでもあれば Yes を出力し、そうでなければ No を出力します。

コード

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

fn main() {
    // 入力
    input! {
        n: usize, // 文字列の長さ
        t: Chars, // 高橋君のリスト
        a: Chars, // 青木君のリスト
    }

    // 両方が'o'の列を探す
    let mut found = false;
    for i in 0..n {
        if t[i] == 'o' && a[i] == 'o' {
            found = true;
            break;
        }
    }

    // 結果を出力
    println!("{}", if found { "Yes" } else { "No" });
}

B 問題

問題

https://atcoder.jp/contests/abc409/tasks/abc409_b

解説

数列 A において以下の条件を満たす x の最大値を求める問題です。

A に、x 以上の要素が重複を含めて x 回以上現れる。

解法としては、x0 から N まで試し、条件を満たすかを判定します。条件を満たす x の中で最大の値が答えとなります。具体的には、x を仮定した際に、数列 A の中で x 以上の要素の個数をカウントします。その個数が x 以上であれば条件を満たします。

コード

abc409b.rs
use proconio::input;

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

    let mut ans = 0;

    // x を 0 から n まで仮定
    for x in 0..=n {
        // 数列 A において x 以上の要素の個数をカウント
        let cnt = count_x_or_more(x, &a);

        // 条件を満たす場合、答えを更新
        if cnt >= x {
            ans = x;
        }
    }

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

// 数列 A において x 以上の要素の個数をカウントする関数
fn count_x_or_more(x: usize, a: &Vec<usize>) -> usize {
    let mut cnt = 0;
    for &val in a {
        if val >= x {
            cnt += 1;
        }
    }
    cnt
}

C 問題

問題

https://atcoder.jp/contests/abc409/tasks/abc409_c

解説

円周上にある点から3点を選び、それらが正三角形を形成する組み合わせの個数を求める問題です。

まず、円周の長さ L が3の倍数でなければ、正三角形を作ることはできません。なぜなら、正三角形を形成するためには、3点間の距離が等しくなる必要があり、その距離は L/3 になります。

次に、円周上の各点の位置を計算します。これは累積和を用いて各点の位置を求めます。各点は始点を0として、隣接距離 d[i] を順に足し合わせて各点の円周上の座標を求めます。その後、各位置に点が何個あるかをカウントします。

最後に、等間隔に L/3 ずつ離れた3点を選ぶ組み合わせを数え上げます。これにより、正三角形を形成する組み合わせの個数を求めることができます。

コード

abc409c.rs
use proconio::input;

fn main() {
    // 入力
    input! {
        n: usize, // 円周上の点の個数
        l: usize, // 円周の長さ
        d: [usize; n - 1], // 各点間の距離
    }

    // 円周の長さが3の倍数でない場合、正三角形は作れない
    if l % 3 != 0 {
        println!("0");
        return;
    }

    // 各点の位置を計算
    let pos_list = calc_cycle_pos(n, l, &d);

    // 各位置に点が何個あるかをカウント
    let cnt_list = get_pos_cnt(l, &pos_list);

    // 正三角形の個数を計算
    let ans = calc_equilateral_triangle_cnt(l, &cnt_list);

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

// 各点の位置を計算
fn calc_cycle_pos(n: usize, l: usize, d: &Vec<usize>) -> Vec<usize> {
    let mut acc_d = vec![0; n];
    for i in 0..n - 1 {
        acc_d[i + 1] = (acc_d[i] + d[i]) % l;
    }
    acc_d
}

// 各位置に点が何個あるかをカウント
fn get_pos_cnt(l: usize, pos_list: &Vec<usize>) -> Vec<usize> {
    let mut cnt = vec![0; l];
    for &pos in pos_list {
        cnt[pos] += 1;
    }
    cnt
}

// 正三角形の個数を計算
fn calc_equilateral_triangle_cnt(l: usize, cnt_list: &Vec<usize>) -> usize {
    let mut ret = 0;
    let sz = l / 3;

    // 等間隔に3つの頂点を選ぶ組み合わせを計算
    for i in 0..sz {
        ret += cnt_list[i] * cnt_list[i + sz] * cnt_list[i + sz * 2];
    }
    ret
}

D 問題

問題

https://atcoder.jp/contests/abc409/tasks/abc409_d

解説

文字列 S について、任意の文字を1個選んで後ろに移動(移動しない場合も含む)して新たにできる文字列 S' の中で、辞書順最小となる文字列を出力する問題です。
文字を1つ移動させる際、①「移動させる文字の特定」、②「挿入位置の特定」の2つを考える必要がありますが、どちらも先頭から順番に調べて特定することが可能です。

① 移動させる文字の特定
文字列 S の先頭から順に、隣り合う文字を比較します。 i 番目の文字 S[i] より後ろの文字 S[i+1] の方が辞書順で小さい場合、S[i] を移動させることで辞書順を小さくできます。候補が複数ある場合は、最初に見つけた文字を移動させた方がより辞書順を小さくできるので、最初に条件を満たす文字を選びます。条件を満たす文字がない場合、元の文字列をそのまま出力します。

② 挿入位置の特定
移動させる文字 S[i] を、i+1 番目以降の文字列の中で辞書順で大きい文字の直前に挿入します。候補が複数ある場合は、最初に見つけた位置に移動させた方がより辞書順を小さくできるので、最初に条件を満たす位置を選びます。挿入位置が見つからない場合は、文字列の末尾に移動させます。

最後に、①②で特定した移動元の文字と挿入位置を使い、新しい文字列を構築します。

コード

abc409d.rs
use itertools::Itertools;
use proconio::{input, marker::Chars};

fn main() {
    // テストケース数の入力
    input! {
        t: usize,
    }
    for _ in 0..t {
        solve();
    }
}

fn solve() {
    // 入力
    input! {
        n: usize, // 文字列の長さ
        s: Chars, // 文字列
    }

    // 移動する文字の位置を特定する
    let move_pos = find_move_pos(n, &s);

    // 入れ替え不要の場合
    if move_pos == n {
        println!("{}", s.iter().join(""));
        return;
    }

    // 移動する文字の挿入位置を特定する
    let insert_pos = find_insert_pos(n, &s, move_pos);

    // 新しい文字列を構築する
    let ans_str = make_str(n, &s, move_pos, insert_pos);

    // 結果を出力
    println!("{}", ans_str.iter().join(""));
}

// 移動する文字の位置を特定する関数
fn find_move_pos(n: usize, s: &Vec<char>) -> usize {
    for i in 0..n - 1 {
        if s[i] > s[i + 1] {
            // 入れ替えるべき文字の位置
            return i;
        }
    }
    // 入れ替え不要
    n
}

// 挿入位置を特定する関数
fn find_insert_pos(n: usize, s: &Vec<char>, move_pos: usize) -> usize {
    for i in move_pos + 1..n {
        if s[move_pos] < s[i] {
            // 挿入すべき位置
            return i;
        }
    }
    // 挿入位置が見つからない場合は末尾
    n
}

// 新しい文字列を構築する関数
fn make_str(n: usize, s: &Vec<char>, move_pos: usize, insert_pos: usize) -> Vec<char> {
    let mut result = Vec::new();

    // 移動箇所より前の文字
    for i in 0..move_pos {
        result.push(s[i]);
    }
    // 移動箇所の後から挿入位置の前までの文字
    for i in move_pos + 1..insert_pos {
        result.push(s[i]);
    }
    // 移動する文字
    result.push(s[move_pos]);
    // 挿入位置以降の文字
    for i in insert_pos..n {
        result.push(s[i]);
    }

    result
}

E 問題

問題

https://atcoder.jp/contests/abc409/tasks/abc409_e

解説

この問題は、木構造上で電子の移動に伴うエネルギー消費を最小化する問題です。
木構造における動的計画法(木DP)を用いて解くことができます。解法としては、以下1~4の手順で解くことができます。

  1. 木の根を1つの頂点に固定します(ここでは頂点0を根とします)。
  2. 再帰的に子ノードを探索し、各ノードで以下を計算します。
    • 子ノードから受け取った電子の合計
    • 子ノードからのエネルギー消費の合計
  3. 各ノードで計算した結果を親ノードに返します。
  4. 最終的に根ノードで計算されたエネルギー消費が、全体の最小エネルギー消費となります。

コード

abc409e.rs
use proconio::{input, marker::Usize1};
const INF: usize = 1 << 60;

fn main() {
    // 入力
    input! {
        n: usize, // 頂点数
        d: [i64; n], // 各頂点にある電子の量
        uvw: [(Usize1, Usize1, i64); n-1] // 辺情報 (u, v, 重み)
    }

    // 無向グラフ
    let mut graph = vec![Vec::new(); n];
    for &(u, v, w) in &uvw {
        graph[u].push((w, v));
        graph[v].push((w, u));
    }

    // 頂点0を根として木DPを実行
    let root = tree_dp(0, INF, &graph, &d);

    // 結果を出力
    println!("{}", root.cost);
}

// 頂点に関する情報を保持する構造体
struct NodeInfo {
    elect: i64, // 頂点が持つ電子の合計
    cost: i64,  // 必要エネルギーの総コスト
}

// 木DPの実装
fn tree_dp(
    cur: usize, // 現在の頂点
    parent: usize, // 親の頂点
    graph: &Vec<Vec<(i64, usize)>>, // 無向グラフ
    d: &Vec<i64>, // 各頂点にある電子の量
) -> NodeInfo {
    // 現在の頂点の電子量とコストを初期化
    let mut cur_info = NodeInfo { elect: d[cur], cost: 0 };

    // 子ノードを探索
    for &(weight, next) in &graph[cur] {
        // 親ノードはみない
        if next == parent {
            continue;
        }

        // 子ノードの情報を取得
        let child_info = tree_dp(next, cur, graph, d);

        // 電子量を合算
        cur_info.elect += child_info.elect;

        // コストを計算
        cur_info.cost += child_info.elect.abs() * weight + child_info.cost;
    }

    // 現在の頂点の情報を返す
    cur_info
}

Discussion