ABC216 を Rust で解く (F問題まで)

はじめに

AtCoder Beginner Contest 216 の F問題までを Rust で解きました！
詳しめに解説をしていこうと思います。
役にたちましたら LGTM いただけると嬉しいです。
AC 確認はしてますが、もっとスマートな書き方がありましたら教えてください。

なぜやるのか

コンテストから時間をあけて解きなおすことで、記憶へ定着すればいいなと思っています。
また、会社で Rust を使うのですが私はまだ慣れておらず、その練習のためにやっています。
(コンテスト本番には C++ を使って参加している水色コーダーです。)

目次

A - Signed Difficulty (Difficulty 17)
B - Same Name (Difficulty 38)
C - Many Balls (Difficulty 145)
D - Pair of Balls (Difficulty 1039)
E - Amusement Park (Difficulty 1084)
F - Max Sum Counting (Difficulty 1541)

前提

この記事にあるコードはこちらを import しています。

# [allow(unused_imports)]
use proconio::{input, fastout, marker::Chars};
# [allow(unused_imports)]
use std::collections::{HashSet, HashMap, BTreeSet, VecDeque};
# [allow(unused_imports)]
use std::cmp::{max, min};
# [allow(unused_imports)]
use itertools::Itertools;

A - Signed Difficulty

文字列処理の問題ですね。

言われていることは簡単ですが、言語によっては処理が難しいです。また、入力値が小数なのですが、 f32 で受け取ったりすると、丸め誤差で数値がずれます。注意です。

私は、文字列で受け取った後に . の前と後ろで分けて、それぞれ数値にして処理しました。

fn main() {
  input! {
    num: String
  }
  let v: Vec<i32> = num
    .split('.')
    .map(|x| x.parse().unwrap())
    .collect_vec();
  
  if v[1] <= 2 { println!("{}-", v[0]) }
  else if 7 <= v[1] { println!("{}+", v[0]) }
  else { println!("{}", v[0]) }
}

B - Same Name

同姓同名の組が存在するか、つまり、重複があるかを判定する問題ですね。

入力は姓と名が空白区切りの文字列でもらえるので、連結します。そのあと、重複除去をします。
・何かの要素が除去された場合は重複があったということで、 Yes を出力
・要素が除去されない場合は重複がなかったということで、No を出力
という判定をすれば答えになります。
要素が除去されたかどうかは、重複除去の処理をする前と後で vector のサイズが変わっているかどうかでわかります。

また、Rust では sort() をして dedup() をすることで重複除去ができます。

fn main() {
  input! {
    n: usize,
    st: [(String, String); n]
  }
  let mut st = st
    .into_iter()
    .map(|(s, t)| s + " " + &t)
    .collect_vec();
    
  st.sort_unstable();
  st.dedup();

  let ans = if st.len() == n { "No" } else { "Yes" };
  println!("{}", ans);
}

C - Many Balls

魔法でボールを増やす問題ですね。

使える魔法は
・1つ増やす
・2倍にする
です。

この「2倍にする」を 1 に対して 60回実行すると、制約の最大値である 10^18 は超えます。「2倍にする」1回ごとに「1つ増やす」を0回か1回やることで、 10^18 までの任意の整数を作れます。これは、数字を作る過程を2進数でやってみるとわかります。
これで問題文にある 120回の回数制限はクリアできました。

また、この問題のコツは、$0$ から $N$ を作るのではなく、 $N$ から $0$ を作ることです。問題文とは逆に考えます。
すると、使う魔法は
・1つ減らす
・2で割る
になります。

あとは、
・N が 奇数の時は 「1つ減らす」
・N が 偶数の時は 「2で割る」
とすれば答えがわかります。
これで出る答えは 問題の答えとして求められているものとは逆になっているので、元に戻す処理が必要なことに注意です。

以上は繰り返し二乗法 (高速累乗) のアルゴリズムとよく似た動きですので、知っていたら解きやすかったかと思います。

fn main() {
  input! {
    mut n: i64
  }
  let mut ans = String::new();
  while n > 1 {
    if n % 2 == 1 {
      ans += "A";
    }
    ans += "B";
    n /= 2;
  }
  ans += "A";
  let mut reversed = String::new();
  for c in ans.chars() {
    reversed = String::from(c.to_string() + &reversed);
  }
  println!("{}", reversed);
}

D - Pair of Balls

公式解説にあるシミュレーションで解いてもいいのですが、少し複雑になってしまいます。
ここでは、グラフを使った解き方を紹介します。

筒からボールを取り出していく動作を考えます。ボールを取り出す "順序" を考えると、ある制約があります。
それは「ボールをとるには一つ上のボールが取られている必要がある」です。

気付くには経験が必要かもしれませんが、順序がある場合はグラフに帰着して解けるものがあります。
順序に関するのはこの制約だけなので、これをグラフにします。つまり、各ボールを頂点として、上のボールから下のボールへの有向辺を張ります。これは有向グラフになります。

作った有向グラフをもとに問題を見ます。「すべてのボールを取り出せる」「取り出せないボールがある」時をそれぞれ考えると、「取り出せないボールがある」時は、グラフが巡回しています。巡回があるとデッドロックのような状態になって、取り出せないボールがあるのもわかります。
つまり、有向グラフの DAG 判定に 帰着できました。ここで、DAG (Directed-Acyclic-Graph) とは「有向非巡回グラフ」のことです。
(脱線しますが、グラフが DAG だとトポロジカルソートという有用なアルゴリズムが使えたりします。)

さて、問題をグラフの DAG 判定に 帰着できたので、これを解きましょう。
調べると、 Rust には petgraph というグラフ関連のアルゴリズム・データ構造がまとまっている crate がありました。
このなかに DiGraph と is_cyclic_directed というものがあったので、これを使って解きます。

petgraph を使わない場合は、訪れた頂点を記録しながら DFS をすると、巡回路の有無を判定できます。

use petgraph::{algo::is_cyclic_directed, graph::DiGraph};

fn main() {
  input! {
    _: usize,
    m: usize,
  }
  let mut edge: Vec<(usize, usize)> = Vec::new();
  for _ in 0..m {
    input!{
      k: usize,
      a: [usize; k],
    }
    for i in 0..k - 1 {
      edge.push((a[i], a[i + 1]));
    }
  }
  let graph = DiGraph::<(), (), usize>::from_edges(&edge);
  let ans = if is_cyclic_directed(&graph) { "No" } else { "Yes" };
  println!("{}", ans);
}

E - Amusement Park

決め打ち二分探索の問題ですね。

これも D問題同様、気付くには経験が必要かもしれません。この問題は二分探索を使わなくてもシミュレーションで解けるので、余計に気付きにくいです。シミュレーションは複雑になるので、二分探索で解きたいところです。

今回のように、 $O(N)$ 以内で1回分の "スコア" を判定できる場合は決め打ち二分探索で解けることが多いです。この "スコア" をどのように設定するかが問題を解くうえで大事なところになってきます。うまく 単調増加 / 単調減少 の関数に帰着させましょう。

今回は、 "スコア" を「アトラクションがとりうる満足度のうち、 arg 以上の個数」 としました。同じアトラクションでも満足度が変われば複数回カウントされることに注意です。この "スコア" が k 以上かどうかを判定します。
これで判定関数ができたので、二分探索をします。

二分探索で境界がわかったら、総和を計算します。
let diff = max(0, a - wa + 1);
ans += diff * (2 * wa + diff - 1) / 2;
の部分です。
各アトラクションごとに、乗る回数を diff に保存して、diff 回乗ったときの満足度の合計を、等差数列の和の公式を使って計算します。

ただ、今回はそれだけでは終わりません。
複数のアトラクションの満足度が同じように減っていくので、これを同時に扱うために、アトラクションに乗った回数が2ずつとか、3ずつとかいう増え方をします。つまり、ぴったり k 回で止まらずに余剰の部分が出てしまいます。
これをうまく処理する必要があって、コードでは cnt に乗った回数を記録しておいて、最後に残り回数分を足しています。
cnt += diff; と ans += (k - cnt) * ac as i64; の部分です。

二分探索の流派について、私はめぐる式の二分探索を、変数名を変えて使っています。解説放送で他のユーザーが紹介した方法で、すぬけ式？と呼ばれているのでしょうか。ここはあまり詳しくありません。。

fn main() {
  input! {
    n: i64,
    k: i64,
    a: [i64; n],
  }

  // arg 以上の要素の個数が k 個以上あるかどうか
  let f = |arg: i64| -> bool {
    let cnt: i64 = a
      .iter()
      .map(|a| max(0, *a as i64 - arg + 1))
      .sum();
    return cnt >= k as i64;
  };

  let mut ac: i64 = 0;
  let mut wa: i64 = 1001001001001;
  while wa - ac > 1 {
    let wj: i64 = (wa + ac) / 2;
    if f(wj) { ac = wj; } else { wa = wj; }
  }

  let mut ans: i64 = 0;
  let mut cnt: i64 = 0;
  for a in a.into_iter() {
    let diff = max(0, a - wa + 1);
    ans += diff * (2 * wa + diff - 1) / 2;
    cnt += diff;
  }
  ans += (k - cnt) * ac as i64;
  println!("{}", ans);
}

F - Max Sum Counting

DP の問題です。

複雑ですが、ひとつずつ読み解いていきましょう。まず、組み合わせを求める問題ですね。取り得る組み合わせは全部で $2^N - 1$ 通りあります。
$max_{i \in S}A_i$ は最大値なので、集合 $S$ に $A_i$ 以下の要素が増えても問題ありません。
$\sum_{i \in S}B_i$ はそうはいかず、集合 $S$ に要素が増えると値も変わってしまいます。
つまり、 $A_i$ のほうが変更にルーズです。これをうまく使いましょう。また、 $A_i$ は最大値なので、ソートをうまく使えそうです。
$B_i$ の和は最大で 5000 * 5000 = 25000000 になるのですが、問題設定より、用があるのは 5000 以下の場合だけです。
ここまで考えると、 dp[i][j] = 0..i を見て、和が j になる組み合わせ というDPテーブルを思いつきます。

要素を A_i の値でソートすることで、$max_{i \in S}A_i$ についてはあまり気にしなくていいようになっています。
特定の順番で処理することで単純になるのは、 ABC214-D や ABC210-D にも似ていますね。

さて、このDPテーブル dp[i][j] = 0..i を見て、和が j になる組み合わせ について、見覚えはありますでしょうか。そう。ナップサックDPですね。細かく見ると違うのですが、テーブルの更新などは似ています。
ナップサックDPはABC032-D が良い練習になります。

ここまでやれば、答えは目前です。
dp[i][j] = 0..i を見て、和が j になる組み合わせ の 0..i について、 i は含まれていないので、それを考慮しながら条件に合うものの和を求めて、答えになります。

fn main() {
  input! {
    n: usize,
    a: [usize; n],
    b: [usize; n],
  }
  let m = 998244353;
  let mut ab: Vec<(usize, usize)> = vec![(0, 0); n];
  for i in 0..n {
    ab[i] = (a[i], b[i]);
  }
  ab.sort_unstable();

  // dp[i][j] = 0..i を見て、和が j になる組み合わせ
  let mut dp = vec![[0_usize; 5010]; n + 1];
  dp[0][0] = 1;
  for i in 1..=n {
    let (_, b) = ab[i - 1];

    for j in 0..=5000 { // knapsack DP
      dp[i][j] = dp[i - 1][j];
      if j >= b {
        dp[i][j] += dp[i - 1][j - b];
        dp[i][j] %= m;
      }
    }
  }

  let mut ans = 0;
  for i in 0..n {
    let (a, b) = ab[i];
    if a < b { continue; }
    for j in 0..=a-b { // b[i] を足した時に、和が a[i] 以下になる組み合わせ
      ans += dp[i][j];
      ans %= m;
    }
  }

  println!("{}", ans);
}

---

最後まで読んでいただきありがとうございます！