👨‍🍳

ABC397:Rustで解く!問題解説

に公開

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

A問題

問題

https://atcoder.jp/contests/abc397/tasks/abc397_a

解説

X の値に応じて、1~3のいずれかを出力します。

  1. X が38.0以上の場合は、1を出力
  2. X が37.5以上の場合は、2を出力
  3. X が37.5未満の場合は、3を出力

コード

ABC397a.rs
use proconio::input;
fn main() {
    // 入力
    input!{
        x: f64,
    }

    // 出力
    if x >= 38.0 {
        println!("1");
    }
    else if x >= 37.5 {
        println!("2");
    }
    else {
        println!("3");
    }
}

B問題

問題

https://atcoder.jp/contests/abc397/tasks/abc397_b

解説

改札機の入退室状態をシミュレーションし、以下の2つのケースに応じて追加の操作回数を求めます。

  1. 次に行う操作が、現在の状態と同じである場合
    → 入室または退室の操作を間に挟む必要があるため、追加の操作が1回必要になります。
  2. 次に行う操作が、現在の状態と異なる場合
    → 次の操作の状態に更新します。

最後まで処理を行った際に入室の状態で終わっていた場合は、退室にするためにさらに追加で1回の操作が必要であることに注意します。

コード

abc397b.rs
use proconio::{input, marker::Chars};
use std::mem::swap;

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

    // 挿入回数
    let mut ans = 0;
    // 現在の状態(true:開いている、false:閉じている)
    let mut cur_st = false;

    // シミュレーション
    for ss in s {
        // 次の状態
        let mut next_st = if ss == 'i' { true } else { false };
        // 同じ状態になっている場合は、操作が必要
        if cur_st == next_st {
            ans += 1;
        }
        // 異なる状態であれば、ステータスを更新
        else {
            swap(&mut cur_st, &mut next_st);
        }
    }
    // 開いている状態で終わっていたら、+1回操作が必要
    if cur_st { 
        ans += 1;
    }

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

C問題

問題

https://atcoder.jp/contests/abc397/tasks/abc397_c

解説

数列 A の分割方法を全て試して、左側と右側に現れる値の種類数の合計の最大値を求めます。
初期状態として、左側には数列 A の最初の要素だけを、右側には残りの N-1 個の要素を置きます。この時、各要素の値と個数を管理するために連想配列を使用します。その後は、右側の要素を一つずつ左側に移動させながら、左側と右側のそれぞれに存在する要素について、値の種類数の合計の最大値を更新していきます。

コード

abc397c.rs
use proconio::input;
use std::collections::HashMap;
use std::cmp::max;

fn main() {
    // 入力
    input!{
        n: usize,
        a: [usize; n],
    }

    // 数列Aを左側と右側に分けて、各要素の値と個数を連想配列で管理
    let mut left_mp = HashMap::new();
    let mut right_mp = HashMap::new();

    // 左側に1個、右側にn-1個ある状態で初期化する
    init_map(&a, n, &mut left_mp, &mut right_mp);

    // 種類数の合計の最大値
    let mut ans = left_mp.len() + right_mp.len();

    for i in 1..n-1 {
        // 右側に含まれているa[i]を左側に移動する
        move_map(&a, i, &mut left_mp, &mut right_mp);
        // 種類数の合計の最大値を更新
        ans = max(ans, left_mp.len() + right_mp.len());
    }

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

fn init_map(a: &Vec<usize>, n: usize, left: &mut HashMap<usize, usize>, right: &mut HashMap<usize, usize>) {
    *left.entry(a[0]).or_insert(0) += 1;
    for i in 1..n {
        *right.entry(a[i]).or_insert(0) += 1;
    }
}

fn move_map(a: &Vec<usize>, idx: usize, left: &mut HashMap<usize, usize>, right: &mut HashMap<usize, usize>) {
    if let Some(&val) = right.get(&a[idx]) {
        // その要素が1個だけの場合、連想配列から削除
        if val == 1 { right.remove(&a[idx]); }
        else {
            *right.entry(a[idx]).or_insert(0) -= 1;
        }
        *left.entry(a[idx]).or_insert(0) += 1;
    }
}

D問題

問題

https://atcoder.jp/contests/abc397/tasks/abc397_d

解説

与えられた式をそのまま使って正整数 (x,y) の解を直接求めるのは難しいため、式を変形して変数の探索範囲を絞り込んで解を探します。まず、元の式を (1) とします。

x^{3} - y^{3} = N (1)

(1) 式を因数分解すると、以下の式に変形できます。

(x - y)(x^{2} + xy + y^{2}) = N (2)

ここで (2) 式について考えると、(x^{2} + xy + y^{2}) は必ず正の値になるため、(x - y) も正の数となり、x > y が成り立ちます。また、(x^{2} + xy + y^{2}) であることから、x^{2}N を超えると等式が成り立たないため、x^{2} \leq N すなわち、x \leq \sqrt{N} となります。さらに、N の最大値は 10^{18} であるため、x \leq \sqrt{N} = 10^{9} となります。よって、(2) 式から、0 < y < x < 10^{9} であることが分かります。
これにより、x,y の上限を絞り込みましたが、x,yの調べる範囲が 1 から 10^{9} と非常に大きいため、さらに条件を絞り込む必要があります。

次に、(2) 式の (x - y) に注目し、(x - y) = d とおくと、x = y + d と置き換えることができます。これを (2) 式に適用すると、以下の式に変形できます。

d(d^{2} + 3dy + 3y^{2}) = N (3)

(3) 式について考えると、左辺はいずれも d の倍数であるため、Nd で必ず割り切れる必要があります。また、(3)の多項式の次数は3であることから、 d^3N を超えると等式が成り立たないため、d^3 \leq N すなわち、d \leq \sqrt[3]{N} = 10^{6} と分かります。
(3) 式から、x - y の差 (d)10^{6} 以内、かつ Nd で割り切れるという条件が分かります。

(1)~(3)より、 d の範囲が 1 から 10^{6}yの範囲が 1 から 10^{9} と分かるので、dの範囲を全て試すことで解を求めることができます。ここで y についてですが、y を大きくすると左辺の値が必ず大きくなり、以下の関係が成り立ちます。
y が小さい場合、左辺 < 右辺
② 等式を満たす y の場合、左辺 = 右辺
y が大きい場合、左辺 > 右辺

そのため、二分探索を用いて、②を満たす y が存在するかを調べれば良いです。
計算量としては、O(\sqrt[3]{N} \log(\sqrt{N})) \approx 10^{8} 程度であるため、実行制限時間内に解くことができます。

コード

abc397d.rs
use proconio::input;

fn main() {
    // 入力
    input! {
        n: i64,
    }

    // x^3 - y^3 = N (1): x > y > 0
    // 式変形すると
    // (x - y)(x^2 + xy + y^2) = N (2) : x^2 < N → x < N^{1/2}
    // x - y = d とおくと、x = y + d
    // 式変形すると
    // (y + d - d)((y + d)^2 + (y + d)y + y^2) = N
    // d*(d^2 + 3*dy + 3*y^2) = N
    // d^3 + 3*d^2*y + 3*d*y^2 = N (3) : d^3 < N → d < N^{1/3}
    // 両辺をdで割る
    // (d^2 + 3*dy + 3*y^2) = N/d (4)
    // (4)について二分探索
    // yが小さい場合、左辺がN/dより小さい
    // 解となるyの場合、左辺がN/dと一致
    // yが大きい場合、左辺がN/dより大きい

    // x-yの範囲
    let d_range = 1_000_000 as i64;

    // 1~d_rangeを全探索
    for d in 1..=d_range {
        // Nはdで割り切れる必要がある
        if n % d != 0 { continue; }
        // 右辺
        let right_side = n / d;    

        // 二分探索で、right(N/d)以下を満たす最大のyを見つける
        let y = binary_search(right_side, d);

        // 左辺と右辺がイコールとなる場合は出力
        if calc_left_side(d, y) == right_side {
            println!("{} {}", y + d, y);
            return;
        }
    }
    // 見つからないので、-1を出力
    println!("-1");
}

fn binary_search(right_side: i64, d: i64) -> i64 {
    let mut ok = 1 as i64;
    let mut ng = 1_000_000_000 + 1;
    for _ in 0..50 {
        let mid = (ok + ng) / 2;
        if calc_left_side(d, mid) > right_side { 
            ng = mid; 
        } else { 
            ok = mid; 
        }
    }
    ok
}

// (d^2 + 3*d*y + 3*y^2)
fn calc_left_side(d: i64, y: i64) -> i64 {
    let mut ret = d * d;
    ret += 3 * d * y;
    ret += 3 * y * y;
    ret
}

E問題

問題

https://atcoder.jp/contests/abc397/tasks/abc397_e

解説

木の根(親)を固定し、葉(子)から順に長さ K のパス(※)が構成できるかを判定する方法を考えます。
(※パス:連続した頂点が辺で結ばれ、各頂点の次数が2以下で閉路を含まない経路)
ここでは1番目の頂点を根とし、各頂点からなる枝の数(次数)を子の枝と親の枝に注目して調べていきます。

(1)子の枝について

  • 各子の枝の長さがKの場合、その枝を取得せず枝を切り離します。
  • 各子の枝の長さがKでない場合、その枝を取得して繋げます。
    このとき、取得した枝のサイズの合計と、取得した枝の数の合計を数えておきます。

(2)親の枝について

  • 枝の長さの合計がKの場合、親の枝を取得せず枝を切り離します。
  • 枝の長さの合計がKでない場合、親の枝を取得して繋げます。

(1)および(2)の結果、取得した枝の数が2以下であればパスの条件を満たしますが、3以上の場合はパスにはなりません。これをDFS(深さ優先探索)を用いて葉から順に全ての頂点について調べ、パスの条件を満たしているかかどうかを判定します。

コード

abc397e.rs
use proconio::{input, marker::Usize1}; 
const INF: usize = 1 << 60;
fn main() {
    // 入力
    input!{
        n: usize, k: usize,
        uv: [(Usize1, Usize1); n*k-1], // インデックスを0始まりにする
    }

    // 無向グラフの辺
    let mut edge = vec![Vec::new(); n*k];
    for (u, v) in uv {
        edge[u].push(v);
        edge[v].push(u);
    }

    // 頂点0を根として、葉から順番に調べる
    let mut f = true;
    dfs(0, INF, &edge, k, &mut f);

    // 出力
    println!("{}", if f {"Yes"} else {"No"});
}

fn dfs(
    cpos: usize, ppos: usize, // 現在の位置、ひとつ前の位置
    edge: &Vec<Vec<usize>>, // 無向グラフの辺
    k: usize, // 切断するパスの長さ
    ans: &mut bool,
) -> usize {
    // 次数(取得した枝の数)
    let mut cnt = 0;
    // パスの長さ(初期パスの長さは1(自身を含む))
    let mut path_length = 1;

    // 子を探索
    for &npos in &edge[cpos] {
        // 親への探索はしない
        if npos == ppos { continue; }
        // 子の枝の長さを取得
        let ret_length = dfs(npos, cpos, edge, k, ans);
        // もらう子の枝の長さが0でなければ、自身と子を繋げる
        if ret_length != 0 { 
            cnt += 1;
            path_length += ret_length;
        }
    }
    // 自身の枝の長さがちょうどKでなければ、自身と親を繋げる
    if cpos != 0 && path_length != k {
        cnt += 1;
    }

    // 次数が3以上ならパスにならないので不可
    if cnt >= 3 {
        *ans = false;
    }

    // パスの長さを返す(Kで割り切れる場合は0とする)
    path_length % k
}

Discussion