👨‍🍳

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

に公開

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

ABC417-A

問題

https://atcoder.jp/contests/abc417/tasks/abc417_a

文字列 S について、先頭から A 文字、末尾から B 文字を取り除いた文字列を出力する問題です。

解説

元の文字列から削除後に残る部分の開始位置と終了位置を計算します。
0-indexの開始位置は A 番目 、終了位置は N - B 番目 となるので、この範囲内の文字を順に出力することで、求める文字列を得ることができます。

コード

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

fn main() {
    // 入力
    input! {
        n: usize, // 文字数
        a: usize, // 削除する先頭文字数
        b: usize, // 削除する末尾文字数
        s: Chars, // 文字列
    }

    // 削除後の読み出し位置(from~to)を求める
    let from = a;
    let to = n - b;

    // [from~to)にある文字を出力
    for i in from..to {
        print!("{}", s[i]);
    }
    println!();
}

ABC417-B

問題

https://atcoder.jp/contests/abc417/tasks/abc417_b

整数列 A から、削除する値が含まれる整数列 B の要素を順に削除し、最終的に残った値を出力する問題です。

解説

整数列 A の中から、整数列 B の各要素を順に削除します。削除は、A の中で B の要素と一致する最初の値を見つけ、それを削除済みの値(ここでは -1)に置き換えることで実現します。
B の全ての要素について操作を行った後、A の中で -1 ではない値を出力します。

コード

abc417b.rs
use proconio::input;

fn main() {
    // 入力
    input! {
        n: usize, // 整数列Aの長さ
        m: usize, // 操作回数
        mut a: [i64; n], // 整数列A
        b: [i64; m], // 削除する値
    }

    // 各操作を実施
    for i in 0..m {
        // 整数列を全探索
        for j in 0..n {
            // 操作する値について、初めて一致したものを削除済みの値(-1)にする
            if b[i] == a[j] {
                a[j] = -1;
                break;
            }
        }
    }

    // 削除済みの値でないものを出力
    for i in 0..n {
        if a[i] != -1 {
            print!("{} ", a[i]);
        }
    }
    println!();
}

ABC417-C

問題

https://atcoder.jp/contests/abc417/tasks/abc417_c

整数列 A に対して、以下の式を満たす組 (i, j) (i < j) の個数を求める問題です。

j - i = A_i + A_j

解説

この式を変形すると、

i + A_i = j - A_j

となります。
この変形により、左辺 i + A_i と右辺 j - A_j が一致するかどうかを確認すれば良いことがわかります。
具体的には、i を固定して、i+1 番目以降にある値と一致する個数を数えます。
i+1 番目以降にある値の個数については、連想配列(HashMap)を用いると効率よく求めることができます。

コード

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

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

    // j - A[j] の分布を計算
    let mut mp = HashMap::new();
    for j in 0..n {
        *mp.entry(j as i64 - a[j]).or_insert(0) += 1;
    }

    let mut ans = 0usize;
    
    // 数列を全探索
    for i in 0..n {
        // i 番目の要素を捨てる
        if let Some(cnt) = mp.get_mut(&(i as i64 - a[i])) {
            *cnt -= 1;
            if *cnt == 0 {
                mp.remove(&(i as i64 - a[i]));
            }
        }
        
        // i+1 番目以降の値について、 i + A[i] と一致する個数を加算する
        if let Some(&cnt) = mp.get(&(a[i] + i as i64)) {
            ans += cnt;
        }
    }

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

ABC417-D

問題

https://atcoder.jp/contests/abc417/tasks/abc417_d

初期テンションが与えられた状態で、N 個のプレゼントを順番にもらった後の最終的なテンションを求める問題です。各プレゼントにはテンションの閾値、増加値、減少値が設定されており、テンションが閾値以下の場合は増加値が加算され、閾値を超える場合は減少値が減算されます。

解説

各プレゼントをもらったときに発生するテンションの閾値 (P) 、増加値 (A) 、減少量 (B) の範囲が小さいことに注目して解きます。これは、動的計画法(DP)を用いて、各プレゼントをもらった後のテンションの状態をDPテーブルに記録し、後ろから順に計算していきます。
DPの状態、遷移は以下のように定義します。

  • 状態
    {DP[\text{cur個選択済み}][\text{現在のテンション}(state)] = \text{最終的なテンション}}

  • 遷移
    \text{cur} 番目のプレゼントをもらうとき、

    現在のテンションが閾値 P 以下の場合:

    dp[cur][state] ← dp[cur+1][state + A]

    現在のテンションが閾値 P より大きい場合:

    dp[cur][state] ← dp[cur+1][state - B]

なお、初期テンションが大きい場合、テンションが下がり続ける区間を考慮する必要があります。この区間については累積和を用いて効率的に計算し、二分探索で特定します。

コード

abc417d.rs
use proconio::input;
const SZ: usize = 1000; // テンションの変化の適応範囲(Pmax+Amax)

fn main() {
    // 入力
    input! {
        n: usize, // プレゼントの数
        pab: [(i64, i64, i64); n], // テンションの閾値、増加値、減少値
    }

    // DP[cur個選択済み][現在のテンション(state)] = 最終的なテンション
    let mut dp = vec![vec![0; SZ + 1]; n + 1];
    // DPを初期化(プレゼントをN個もらったときの各状態)
    init_dp(&mut dp, n);
    // DPの各状態を求める
    calc_dp(&mut dp, n, &pab);

    // テンションの減少を連続で選んだ時の累積和を求める
    let acc_b = calc_acc_decrease(n, &pab);

    input! {
        q: usize, // クエリ数
        x: [i64; q], // 各値
    }

    // 各値から始めた時のテンションを求める
    for xx in x {
        // N回連続でテンションが下がっても、SZより大きい場合
        if xx - acc_b[n] > SZ as i64 {
            let ans = xx - acc_b[n];
            println!("{}", ans);
        }
        // 0~N回連続でテンションが下がった後、テンションの増減あり
        else {
            // 連続でテンションが下がる回数を求める
            let cnt = calc_continuous_decrease(n, &acc_b, xx);
            // cnt回テンションが下がった後の値(SZ以下)
            let ans = xx - acc_b[cnt as usize];
            // 残りの増減分は、DPテーブルから調べる
            println!("{}", dp[cnt as usize][ans as usize]);
        }
    }
}

fn init_dp(dp: &mut Vec<Vec<i64>>, n: usize) {
    // N個プレゼントをもらった状態のテンションを初期化
    for state in 0..=SZ {
        dp[n][state] = state as i64;
    }
}

fn calc_dp(dp: &mut Vec<Vec<i64>>, n: usize, pab: &Vec<(i64, i64, i64)>) {
    // 後ろから順に調べる
    for cur in (0..n).rev() {
        for state in 0..=SZ {
            // P以下のときは、Aだけ得る
            if state <= pab[cur].0 as usize {
                let nstate = state as i64 + pab[cur].1;
                dp[cur][state] = dp[cur + 1][nstate as usize];
            }
            // Pを超えているときは、Bだけ失う
            else {
                let nstate = state as i64 - pab[cur].2;
                if nstate < 0 {
                    dp[cur][state] = dp[cur + 1][0];
                } else {
                    dp[cur][state] = dp[cur + 1][nstate as usize];
                }
            }
        }
    }
}

fn calc_acc_decrease(n: usize, pab: &Vec<(i64, i64, i64)>) -> Vec<i64> {
    let mut acc_b = vec![0; n + 1];
    for i in 0..n {
        acc_b[i + 1] = acc_b[i] + pab[i].2;
    }
    acc_b
}

fn calc_continuous_decrease(n: usize, acc_b: &Vec<i64>, x: i64) -> i64 {
    // 二分探索で回数を調べる
    let mut ng = -1;
    let mut ok = (n + 1) as i64;
    for _ in 0..16 {
        let mid = (ng + ok) / 2;
        if x - acc_b[mid as usize] <= SZ as i64 {
            ok = mid;
        } else {
            ng = mid;
        }
    }
    ok
}

ABC417-E

問題

https://atcoder.jp/contests/abc417/tasks/abc417_e

無向グラフにおいて、頂点 X から頂点 Y への経路のうち、辞書順で最小のものを求める問題です。

解説

頂点 X から頂点 Y への経路の探索は、深さ優先探索(DFS)を用いることで全ての経路を調べることができます。今いる頂点から次に移動する頂点を選ぶ際に、まだ探索していない頂点の中で最も小さい頂点番号を選ぶようにします。これにより、辞書順で最小の経路が得られるので、DFS中に頂点 Y に到達した時点で探索を終了し、経路を出力します。
なお、頂点 Y に移動できなかった場合はひとつ前の頂点に戻ります。ひとつ前の頂点に戻った際、その頂点は経路から取り除く必要がありますが、後戻りが発生した頂点以降から頂点 Y へ行く方法がないと分かっているため、探索状況は戻さないことに気をつけます。

コード

abc417e.rs
use itertools::Itertools;
use proconio::input;

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

fn solve() {
    input! {
        n: usize, // 頂点数
        m: usize, // 辺の数
        x: usize, y: usize, // 始点と終点
    }

    // 無向グラフの構築
    let mut graph = vec![Vec::new(); n + 1];
    for _ in 0..m {
        input! {
            u: usize, v: usize,
        }
        graph[u].push(v);
        graph[v].push(u);
    }
    
    // グラフの頂点順を昇順にする
    for i in 0..=n {
        graph[i].sort();
    }

    // 探索済みの頂点 (xのみ探索済み状態で初期化)
    let mut seen = vec![false; n + 1];
    seen[x] = true;

    // DFSの実行
    dfs(x, y, &graph, &mut seen, &mut vec![x], &mut false);
}

fn dfs(
    cur: usize, goal: usize, // 現在の頂点、終点の頂点
    graph: &Vec<Vec<usize>>, // 無向グラフ
    seen: &mut Vec<bool>, // 探索済みの頂点
    route: &mut Vec<usize>, // 始点から終点までの経路
    reached: &mut bool, // 到達済み判定
) {
    // 終点に一度到達したら探索終了
    if *reached {
        return;
    }

    // 終点に到達
    if cur == goal {
        // 到達済みにする
        *reached = true;
        // 経路を出力
        println!("{}", route.iter().join(" "));
        return;
    }

    // 終点に到達していない場合は、次に移動する頂点を探す
    for &next in &graph[cur] {
        // 一度探索した頂点は見ない
        if seen[next] {
            continue;
        }

        // 頂点を探索済みにする
        seen[next] = true;
        // 選択した頂点を経路に追加
        route.push(next);
        // 次の頂点を調べる
        dfs(next, goal, graph, seen, route, reached);
        // 選択した頂点を経路から取り除く
        route.pop();
    }
}

Discussion