👨‍🍳

ABC385:Rustで解く!問題解説

2024/12/25に公開

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

A問題

  • 問題

https://atcoder.jp/contests/abc385/tasks/abc385_a

  • 解説

3つの整数をグループに分ける方法を順番に考えます。

(1) 2つのグループに分けた場合
昇順にソートした値の中で、最初の2つの値の和が、3番目の値と一致する場合に条件を満たします。

(2) 3つのグループに分けた場合
3つの値がすべて同じである場合に条件を満たします。

  • コード
use proconio::input;

fn main() {
    input! {
        mut abc: [usize; 3],
    }
    // 昇順にソート
    abc.sort();

    // (1) 2つのグループに分けた場合
    if abc[0] + abc[1] == abc[2] {
        println!("Yes");
    }
    // (2) 3つのグループに分けた場合
    else if abc[0] == abc[1] && abc[1] == abc[2] {
        println!("Yes");
    } else {
        println!("No");
    }
}

B問題

  • 問題

https://atcoder.jp/contests/abc385/tasks/abc385_b

  • 解説

スタート地点(X, Y)から文字列TLRUDに従い、移動可能であれば移動し、到達した家のマス@の数を数えます。なお、一度着いた家のマスは重複して数えないようにするため、到達した家のマスを.に置き換えます。

  • コード
use proconio::{input, marker::Usize1, marker::Chars};
fn main() {
    input!{
        h: usize,                 // グリッドの高さ
        w: usize,                 // グリッドの幅
        x: Usize1,                // スタート地点のX座標
        y: Usize1,                // スタート地点のY座標
        mut s: [Chars; h],       // グリッドの状態
        t: Chars,                // 移動指示の文字列
    }
    // 上左下右の移動
    const D4 : [(usize, usize); 4] = [(!0, 0), (0, !0), (1, 0), (0, 1)];

    // 到達した家の数
    let mut cnt = 0;           
    // 今の座標
    let mut cpos = (x, y); 
    // 移動後の座標
    for &tt in &t {
        let mut npos = (cpos.0, cpos.1); 
        let di;
        
        // 移動指示に応じた方向を決定
        if tt == 'U' { di = 0; }
        else if tt == 'L' { di = 1; }
        else if tt == 'D' { di = 2; }
        else { di = 3; }
        
        // 新しい座標を計算
        npos.0 = cpos.0.wrapping_add(D4[di].0);
        npos.1 = cpos.1.wrapping_add(D4[di].1);
        
        // 移動可能かどうかをチェック
        if npos.0 >= h || npos.1 >= w || s[npos.0][npos.1] == '#' {
            continue;
        }
        
        // 家に到達した場合
        if s[npos.0][npos.1] == '@' {
            cnt += 1;
            s[npos.0][npos.1] = '.';
        }
        cpos = npos;
    }
    // 座標と到達した家の数を出力
    println!("{} {} {}", cpos.0 + 1, cpos.1 + 1, cnt);
}

C問題

  • 問題

https://atcoder.jp/contests/abc385/tasks/abc385_c

  • 解説

結論から言うと、この問題は全探索で解くことができます。まず、始めに選ぶ位置を固定し、指定した個数だけ先を見て同じ値であれば、そのまま続けていきます。1個先を見た場合は最大でN回見ることになりますが、2個先を見た場合は\frac{N}{2}回、3個先を見た場合は\frac{N}{3}回と、見る回数は徐々に減少していきます。

この回数は調和級数で計算することが出来、

\begin{aligned} \frac{N}{1} + \frac{N}{2} + ... + \frac{N}{n} &= N \times \int_1^n \frac{1}{n} dn \\ &= N \log{N} \end{aligned}

となります。これを全ての位置で試すと上の計算結果のN倍、計算量はO(N^2 \log{N})になります。Nの制約は3000以下のため、O(N^2 \log{N})でも実行時間に間に合います。

  • コード
use proconio::input;
use std::cmp::max;

fn main() {
    input! {
        n: usize,
        h: [usize; n],
    }
    // 自身は必ず選ぶので、答えは1以上
    let mut ans = 1;
    
    // 始めに選ぶ位置
    for spos in 0..n {
        // 選んだ値
        let now = h[spos];
        // k個ずつ進める
        for k in 1..n {
            // 選んだ位置からの同じ値の数
            let mut cnt = 1;
            // 現在の位置
            let mut cpos = spos;
            // k個先が同じ値であれば、続けて選ぶ
            while cpos + k < n && h[cpos + k] == now {
                cpos += k;
                cnt += 1;
            }
            ans = max(ans, cnt);
        }
    }
    println!("{}", ans);
}

D問題

  • 問題

https://atcoder.jp/contests/abc385/tasks/abc385_d

  • 解説

B問題と同様に、スタート地点(X, Y)からLRUDの指示に従ってCだけ進む操作を繰り返します。ただし、家の座標の取りうる範囲は-10^910^9と非常に広いため、配列に直接格納することはできません。
移動は縦または横に行うため、縦方向の移動時には横列を管理している連想配列、横方向の移動時には縦列を管理している連想配列を使用します。そして、家の座標を順序付き集合で保持します。

また、Cだけ移動する操作を行った際、家が複数存在する場合、それらを全て順序付き集合から取り出す必要があります。これは順序付き集合の二分探索を用いて、移動前から移動後の間に含まれる座標を調べることで、効率的に特定できます。

  • コード
use proconio::input;
use std::collections::{HashMap, BTreeSet};
use std::mem::swap;

fn main() {
    input! {
        n: usize, m: usize,
        sx: i64, sy: i64,
        mut xy: [(i64, i64); n],
        dc: [(char, i64); m],
    }

    // 家の座標を格納するハッシュマップ
    let mut row_mp: HashMap<i64, BTreeSet<i64>> = HashMap::new();
    let mut col_mp: HashMap<i64, BTreeSet<i64>> = HashMap::new();
    for &(x, y) in &xy {
        col_mp.entry(x).or_insert_with(BTreeSet::new).insert(y);
        row_mp.entry(y).or_insert_with(BTreeSet::new).insert(x);
    }
    
    // スタート地点の座標
    let mut now = (sx, sy);
    // 到達した家の数
    let mut cnt = 0;

    for &(d, c) in &dc {
        // 次の座標
        let mut next = now;
        match d {
            'U' => next.1 += c,
            'D' => next.1 -= c,
            'L' => next.0 -= c,
            'R' => next.0 += c,
            _ => continue,
        }

        // 縦移動の場合
        if d == 'U' || d == 'D' {
            move_range(
                &mut col_mp,
                &mut row_mp,
                now.0,
                now.1,
                next.1,
                &mut cnt,
            );
        // 横移動の場合
        } else {
            move_range(
                &mut row_mp,
                &mut col_mp,
                now.1,
                now.0,
                next.0,
                &mut cnt,
            );
        }
        // 更新された位置を現在位置に設定
        now = next; 
    }

    // 最終的な位置と到達した家の数を出力
    println!("{} {} {}", now.0, now.1, cnt);
}

fn move_range(
    primary_map: &mut HashMap<i64, BTreeSet<i64>>, // 縦または横の連想配列
    secondary_map: &mut HashMap<i64, BTreeSet<i64>>, // 横または縦の連想配列
    primary_key: i64, // 移動対象のキー
    mut start: i64,   // 開始位置
    mut end: i64,     // 終了位置
    cnt: &mut i64,    // 家のカウントを更新するための参照
) {
    // 開始位置と終了位置を入れ替え
    if start > end { 
        swap(&mut start, &mut end)
    } 

    // 遷移対象の連想配列から値を取得
    if let Some(set) = primary_map.get_mut(&primary_key) {
        let mut values_to_remove = Vec::new();

        // 指定された範囲から値を選択
        for &val in set.range(start..) {
            if val > end {
                break;
            }
            // 値を収集
            values_to_remove.push(val);
        }
        
        // 対象の値を削除
        for val in values_to_remove {
            set.remove(&val);
            if let Some(secondary_set) = secondary_map.get_mut(&val) {
                secondary_set.remove(&primary_key);
            }
            // 家の数をカウント
            *cnt += 1;
        }
    }
}

E問題

  • 問題

https://atcoder.jp/contests/abc385/tasks/abc385_e

  • 解説

問題文にある木をどれだけ大きくできるかを考えます。考える必要がある頂点は赤色(木の根)、青色(赤色の頂点から見た子の頂点)、緑色(青色の頂点から見た子の頂点)の3種類がありますが、赤色の頂点を決めると他の色の頂点も自動的に決まるので、赤色から考えることにします。
青色の頂点は1個以上選ぶことができますが、緑色の頂点の個数はすべて同じにする必要があります。
緑色の頂点はなるべく減らさないようにすることを考えると、青色の頂点は連なっている緑色の頂点が多いものを優先して選ぶ方が最適であると考えられます。
まとめると、赤色の頂点を全探索しつつ、青色の頂点と緑色の頂点の数が最も取れる組み合わせを探せばよいです。

  • コード
use proconio::{input, marker::Usize1};
use std::cmp::max;

fn main() {
    input! {
        n: usize,
    }
    // 隣接リスト
    let mut uv = vec![Vec::new(); n];
    // 頂点の次数
    let mut deg = vec![0; n];

    // 辺の情報を読み込む
    for _ in 0..n-1 {
        input! {
            u: Usize1, v: Usize1,
        }
        uv[u].push(v); deg[u] += 1;
        uv[v].push(u); deg[v] += 1;
    }

    // 木の大きさを最大化する
    let mut tree_size = 0;

    // 木の根の頂点を全探索
    for cpos in 0..n {
        // 根と隣接する頂点を全て選択し、次数の降順でソート
        let mut select_deg = Vec::new();
        for &npos in &uv[cpos] {
            select_deg.push(deg[npos]);
        }
        select_deg.sort_unstable_by(|a, b| b.cmp(a)); // 降順でソート

        // 降順でi番目までの頂点を選んで、題意の木を作る
        for i in 0..select_deg.len() {
            // 選んだ頂点数のminの分だけ葉を残す。それをi+1セット、根1個を選択
            let sz = select_deg[i] * (i + 1) + 1;
            if sz <= n {
                tree_size = max(tree_size, sz);
            }
        }
    }
    // 削除する頂点数を求める
    println!("{}", n - tree_size);
}

Discussion