ABC220 を Rust で解く (D問題まで)

はじめに

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

なぜやるのか

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

目次

A - Find Multiple (Difficulty 14)
B - Base K (Difficulty 58)
C - Long Sequence (Difficulty 119)
D - FG operation (Difficulty 664)

前提

この記事にあるコードはこちらを 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 - Find Multiple

b / c * c で、c の倍数の中で b 以下の最大のものがわかります。
この値が a 以上かどうかを判定して答えがわかります。

fn main() {
  input! {
    a: i32,
    b: i32,
    c: i32
  }
  let val = b / c * c;
  let ans = if val < a { -1 } else { val };
  println!("{}", ans);
}

B - Base K

from_str_radix(&str, u32) で、文字列から基数を指定した変換をできます。
変換先は10進数固定なので、 N進数 -> 10進数 の変換ができます。

この関数は今回はじめて知ったので調べてみました。
数値を表す型には実装されていて、変換元は 2進数から36進数に対応しているとのことです。
-> https://doc.rust-lang.org/std/primitive.u64.html#method.from_str_radix
-> https://doc.rust-lang.org/std/?search=from_str_radix

fn main() {
  input! {
    k: u32,
    a: String,
    b: String
  }
  let ans = i64::from_str_radix(&a, k).unwrap() * i64::from_str_radix(&b, k).unwrap();
  println!("{}", ans);
}

C - Long Sequence

問題文をそのまま実装すると $O(X)$ になります。これでは時間内に解けません。
この計算を高速化することを考えましょう。
「$A_i$ の要素が全て足されている部分」は何度も繰り返されるので、ここは削減できそうです。

数列A の和を sum とすると、この部分には 数列Aが x / sum 個分ふくまれますね。
このときの項数は (x / sum) * n で、和は (x / sum) * sum になります。

あとは、これまでの和：(x / sum) * sum に 数列A の要素をひとつずつ足していって、 x を超える地点を調べれば答えがわかります。

類第にABC208-Cがあります。また、この問題が発展すると ABC175-D などになります。

fn main() {
  input! {
    n: i64,
    a: [i64; n],
    mut x: i64
  }
  let sum: i64 = a.iter().sum();
  let mut ans = (x / sum) * n;
  x %= sum;
  for v in a.iter() {
    ans += 1;
    x -= v;
    if x < 0 {
      break;
    }
  }
  println!("{}", ans);
}

D - FG operation

DP の問題ですね。よくある形の DP ですので、解けるようになると DP の理解が進むと思います。

左端からふたつが計算対象です (どこから計算してもいいと誤読していて混乱しました)。

dp[i] = 左端が i の時の組み合わせ数 として解きます。
このようにDPテーブルを設定すると、操作の種類にたいして線形時間で問題を解けます。
つまり、 操作の種類を Q、A の種類を|A| として、全探索すると $O(Q^{N-1})$ となるところを $O(QN|A|)$ で解けるようになります。
今回の問題は $Q=2$ と $2 \le N \le 10^5$ と $|A| = 10$ なので、制限時間内に解けそうですね。

私は dp[i] = 左端が i の時の組み合わせ数 としてメモリを節約しましたが、
dp[i][j] = 前から i まで見て、左端の数字が j の時の組み合わせ数 としたほうが考えやすいかと思います。
メモリ使用量的にも制約の範囲に収まります。

ほか、アルゴリズムの挙動はコメントに書きましたのでご覧ください。

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

  let m = 998244353;

  // dp[i] = 左端が i の時の組み合わせ数
  let mut dp = vec![0; 10];
  dp[a[0]] = 1; // 1番目
  for i in 1..n { // 2番目以降
    let mut next = vec![0; 10];
    for j in 0..10 {
      // 左端が j
      // 左端から2番目の値がa[i]

      // 操作F 足し算
      // 計算結果は (a[i] + j) % 10
      next[(a[i] + j) % 10] += dp[j];
      next[(a[i] + j) % 10] %= m;

     // 操作G 掛け算
      // 計算結果は (a[i] * j) % 10
      next[(a[i] * j) % 10] += dp[j];
      next[(a[i] * j) % 10] %= m;
    }
   dp = next;
  }
  for i in 0..10 {
    println!("{}", dp[i]);
  }
}

---

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