ABC209 を Rust で解く (E問題まで)

はじめに

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

なぜやるのか

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

目次

A - Counting (Difficulty 5)
B - Can you buy them all? (Difficulty 13)
C - Not Equal (Difficulty 285)
D - Collision (Difficulty 686)
E - Shiritori (Difficulty 2153)

前提

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

問題文がすごく単純ですね。
解法も単純で、B - A + 1 が答えになります。
これは具体例を考えると分かると思います。
計算結果が $0$ 未満になることがあるので、if や max などで対処しましょう。

fn main() {
  input! {
    a: i32,
    b: i32
  }
  println!("{}", max(0, b - a + 1));
}

B - Can you buy them all?

偶数番目を値引きしながら合計を計算する問題ですね。
$A$ の偶数番目の要素から $1$ を引いたあと、総和を計算すれば答えがわかります。
この操作には for を使いましょう。
インデックスは $0$ から始まるので、偶数番目のインデックスは奇数になることに注意です。
つまり、 i % 2 == 1 となるものが偶数番目になります。

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

  let mut sum = 0;
  for i in 0..n as usize {
    sum += a[i];
     if i % 2 == 1 {
       sum -= 1;
     }
  }
  let ans = if sum <= x { "Yes" } else { "No" };
  println!("{}", ans);
}

もしくは、イテレータをつかって総和を計算することもできます。
a.iter().sum() ですね。
偶数番目のものの価格が $1$ 引かれるということは、全体でみると、偶数番目の個数と同じ価格が値引きされるということです。$N$ 個の要素のうち、偶数番目のものは N / 2 個になるので、全体の合計から 'N / 2' を引けば答えがわかります。
これは、整数の割り算では小数点以下が切り捨てられることを利用しています。

これを実装すると、以下のようにも書けます。

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

  let sum = a.iter().sum::<usize>() - n / 2;
  let ans = if sum <= x { "Yes" } else { "No" };
  println!("{}", ans);
}

C - Not Equal

これはすこし考えにくい問題だと思いました。
サンプルの場合や、$C$ がすべて $100$ の場合などを試していると、$C$ が昇順だと考えやすいことがわかります。
また、さらに考察を進めると、 $C$ の順番は入れ替えてもいいことがわかります。
なので、昇順でソートしましょう。
「ソートしても答えが変わらないなら、とりあえずソートする」は競プロでよくわれるテクニックです。
一般に、ソートすると考察をしやすくなることが多いのでオススメです。

さて、ここまでわかると、次の考察に進めます。

$C_i$ を見ているときは、それまでに $i - 1$ 個の数字が使われていることがわかります。
たとえば先頭の $C_1$ を見ているときは $0$個で、次の $C_2$ を見ているときは $C_1$ で選んだ $1$ 個の数字が使われています。
また、$C$ は昇順でソートされているので、 それまでに使われている数は全て、 $C_i$ 以下だということがわかります。
$C_i$ 以下の数字が $i - 1$ 個使われているので、 $C_i$ を見ているときに選べる数字は $C_i - (i - 1)$ 個になります。
つまり、$C$ を $1$ つ見るごとに $C_i - (i - 1)$ 通りの組み合わせが増えます。
あとは、これを掛け合わせれば答えがわかります。

以下、注意事項です。
$C_i - (i - 1)$ は負になることもありますが、負になる前に $0$ になる瞬間があります。
$0$ が掛けられると答えが $0$ になり、以降は何を掛けても $0$ のままです。
ここは少し注意ですね。

また、今回は答えがとても大きくなるので、 $10^9 + 7$ で割った余りが必要です。
今回の問題に必要な性質だけ書きますが、掛けるたびに %= 1000000007 をすれば問題ありません。
ただ、$C_i$ が最大で $10^9$ 程度なので、%= 1000000007 を計算する直前は最大で $18$桁になってしまいます。
こうなると、$32$bit 整数の範囲には収まりません。u64 や i64 の出番です。
i64 は $9*10^{18}$ 程度まで表現できて、u64 だと i64 の二倍の範囲まで表現できます。
なので、i64 と u64 のどちらかを使えば、$18$桁の数字も範囲内に収まります。

さて、ここまでわかれば実装です。
長々と書いてしまいましたが、実装はシンプルです。

fn main() {
  input! {
    n: usize,
    mut c: [u64; n]
  }
  c.sort();
  let MOD = 1000000007;
  let mut ans: u64 = 1; // 乗算の単位元である 1 で初期化する
  for i in 0..n as usize {
    ans *= c[i] - i as u64;
    ans %= MOD;
  }
  println!("{}", ans);
}

D - Collision

まずわかるのは、街を頂点、道路を辺としたグラフの問題ということですね。
また、「頂点が $N$個、辺が $N-1$個、連結」という性質があります。これは「木」という名前の付いた特殊なグラフです。
木だということは問題文には明記されていないので、自力で気づく必要があります。
ちなみに、問題文の以下の部分よりわかります。

N個の街と N−1 本の道路からなり
どの街からどの街へもいくつかの道路を通ることで移動できます

さらに、辺に向きがなくて、コストもないので、重みなし無向木だとわかります。
以下の部分よりわかります。

双方向に結んでおり
道路は全て同じ長さです

さて、この重みなし無向木だという前提をもって問題を見ます。
ひとつわかるのは、最短経路はひとつに定まるということですね。(2通りあったりはしないということです。)
つまり、高橋君と青木君は必ずどこかで出会います。
また、実際に試したりすると、 $c_i$ と $d_i$ の最短距離が
・偶数の時は街で出会う
・奇数の時は道路上で出会う
ということがわかります。

つまり、最短距離がわからなくても、距離の偶奇がわかれば解けるようになりました。
ちなみに、ある街からある街への最短距離を求めるには BFS または DFS をすればよくて、これは $O(N)$ の計算量です。
このクエリが $Q$個あるので、全体の計算量は $O(NQ)$ になります。
$1 \le N, Q \le 10^5$ なので、これでは時間内に解けませんね。

なので、何かを速くすることを考えます。上でも書きましたが、「最短距離がわからなくても、距離の偶奇がわかれば」解けます。この方針で考えましょう。

まず、任意の街からスタートして各街への最短距離をだしてみます。BFS か DFS で $O(N)$ で計算できます。そのあとで、クエリである 「$c_i$ から $d_i$ への最短経路」を考えます。

スタートの街が $c_i$ から $d_i$ への最短経路上にあれば、「スタートから $c_i$への距離」 + 「スタートから $d_i$ への距離」が最短距離になります。そのまま偶奇を調べれば答えがわかります。

スタートの街が $c_i$ から $d_i$ の最短経路上にないときは、残念ながら $c_i$ から $d_i$ への最短経路はわかりません。ですが、実は距離の偶奇はわかります。というのも、「$c_i$ から $d_i$ の最短距離」と「スタートの街を経由して $c_i$ から $d_i$ へ行く最短距離」の偶奇が一致するためです。

これについて説明します。
スタートの街を経由するために「余分な距離」がかかってしまいますが、これは同じ道を往復することになります。すると、同じ道を2回通ることになるので、「余分な距離」は偶数になりますね。
「$c_i$ から $d_i$ への最短距離」に偶数を足したものが「スタートの街を経由して $c_i$ から $d_i$ へ行く最短距離」になる。
つまり、偶数の性質から、「$c_i$ から $d_i$ の最短距離」と「スタートの街を経由ながら $c_i$ から $d_i$ へ行く最短距離」の偶奇が一致します。

まとめると、「スタートから $c_i$ への距離」 + 「スタートから $d_i$ への距離」の偶奇を調べることで答えを出せるようになりました！
これだと $O(N)$ でスタートから他の全点への距離を調べた後、各クエリに $O(1)$ で答えられます。
全体では $O(N + Q)$ になり、時間内に解けそうです。

さて、ここまでわかれば実装です。
実装上の注意はコメントに書きました。

fn main() {
  input! {
    n: usize,
    q: usize,
    ab: [(usize, usize); n - 1],
    cd: [(usize, usize); q]
  }
  // vec![初期値; 個数]
  let mut graph = vec![vec![0; 0 as usize]; n];
  for i in 0..n - 1 {
    let (a, b) = ab[i];
    // 無向グラフなので、双方向の辺をつくる
    graph[a - 1].push(b - 1);
    graph[b - 1].push(a - 1);
  }

  let INF = 1 << 30;

  // 0 の街からの距離を dist に保存する
  // 0 の街 (自分自身) への距離は 0 
  let mut dist = vec![INF; n];
  dist[0] = 0;

  // queue をつかって BFS をする
  // もしここで stack をつかえば DFS になる
  let mut que: VecDeque<usize> = VecDeque::new();

  // 0 の街からはじめる
  que.push_back(0);

  // que が空になるまでループする
  while let Some(now) = que.pop_front() {
    // いまは now の街にいて、now の街からいける街を for でループする
    for next in &graph[now] {
      // 借用権にひっかからないようにコピーして、ついでに shadowing する
      let next = *next;

      // すでに来た街かを確認して、来ていたらなにもしない
      if dist[next] != INF { continue; }

      // 距離を更新して、次の街に行く
      dist[next] = dist[now] + 1;
      que.push_back(next);
    }
  }

  for i in 0..q {
    let (c, d) = cd[i];
    // 距離の和 (差でも可) が偶数なら街、奇数なら道路で出会う
    let ans = if (dist[c - 1] + dist[d - 1]) % 2 == 0 { "Town" } else  { "Road" };
    println!("{}", ans);
  }
}

E - Shiritori

(黄色diff なので、これを解く人はわかっていそうな知識は省略します。)

この問題を解くにあたり必要なのは、有向グラフに帰着させるという発想です。
この問題を見るまで気づいていなかったのですが、しりとりって有向グラフに帰着できるんですね。

この有向グラフについて考えると、多重辺・自己ループ・ループありという性質を持ちますね。
この場合、使うアルゴリズムの候補が絞られます。たとえばトポロジカルソートのようなアルゴリズムは使えません。

グラフ、ゲーム、、、などと考えていると、後退解析で解けそうだという思いが沸いてきます。
(EDPC-K を解いていたので思いついた部分もあります。)

なので、まず端点から考えましょう。端点とは遷移先のないノードのことです。
これは遷移先がない (しりとりを繋げられない) ので「負け」のノードです。
次に、端点に遷移できるノードを考えると、「勝ち」のノードになりますね。

これと同様に、ループの場合や、遷移先が複数ある場合を考えます。
すると、以下の条件が出てくるかと思います。
・「負け」のノードに遷移できるノードは「勝ち」
・遷移先のノードが「勝ち」だけの場合は「負け」
・遷移先のノードが「自分」だけ (自己ループ) の場合は「引き分け」
・遷移先のノードが「引き分け」だけの場合は「引き分け」
・遷移先のノードが「引き分け」と「負け」の場合は「勝ち」
・遷移先のノードが「引き分け」と「勝ち」だけの場合は「引き分け」

ここでひとつ大事なのが、[ 遷移先のノードが「勝ち」だけの場合は「負け」 ] を判定するために、元のグラフで言う出次数（反転させたグラフで言う入次数）を記録する必要が出てきました。

ここまでわかれば実装に入ります。

さて実装ですが、アルファベット3文字をノードにします。
ここで、アルファベット3文字をどう表現するかに工夫がいります。
アルファベットは小文字と大文字があるので、全部で52種類です。
これを3文字なので、そのまま52進数で考えてみます。$52*52*52=140608$ なので、空間量的にも問題なさそうですね。
入力で出現した文字だけを考えてもいいのですが、HashMap 等が出てきて少し複雑になりそうです。
空間量的には問題ないので、全通りを試しましょう。

他、必要なことはコメントに書いたつもりです。

fn c_to_usize(c: char) -> usize {
  // a: 97 / A : 65
  // println!("a: {} / A : {}", 'a' as usize, 'A' as usize);
  if c.is_ascii_lowercase() {
    c as usize - 97 + 26
  } else {
    c as usize - 65
  }
}

fn s_to_usize(s: &[char]) -> usize {
  s.iter()
   .map(|&c| c_to_usize(c))
   .fold(0, |digit, c| digit * 52 + c)
}

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

  let MAX_N: usize = 52 * 52 * 52;
  let mut graph_rev = vec![vec![0; 0 as usize]; MAX_N];
  let mut in_degree_rev = vec![0; MAX_N];
  let mut ans = vec![0; MAX_N];
  let mut input_tail_idx = vec![0; n];

  for (i, s) in s.into_iter().enumerate() {
    let head_idx = s_to_usize(&s[0..3]);
    let tail_idx = s_to_usize(&s[s.len() - 3..]);

    // 後退解析をするので、逆向きに辺を張る
    graph_rev[tail_idx].push(head_idx);
    // 入り次数を記録する
    in_degree_rev[head_idx] += 1;
    // 答えを出すときのために値を保存しておく
    input_tail_idx[i] = tail_idx;
  }

  // 0: 不明 or loop 
  // 1: 勝ち
  // 2: 負け
  let undefined = 0;
  let win = 1;
  let lose = 2;

  // 端点を始点にして BFS をする
  let mut que: VecDeque<usize> = VecDeque::new();
  for i in 0..MAX_N {
    if in_degree_rev[i] == 0 {
      // 端点を開始点にする
      // 端点は複数個あるので全探索する
      que.push_back(i);

      // 端点は [負けノード]
      ans[i] = lose;
    }
  }

  while let Some(now) = que.pop_front() {
    for next in &graph_rev[now] {
      let next = *next;
      if ans[next] != undefined { continue; }

      in_degree_rev[next] -= 1;
      if ans[now] == lose {
        // next から [負けノード] に遷移できるなら
        // next は [勝ちノード]
        ans[next] = win;
        que.push_back(next);
      }
      else if in_degree_rev[next] == 0 {
        // next のノードから遷移できるノードが全て勝ちノードなら
        // next は [負けノード]
        ans[next] = lose;
        que.push_back(next);
      }
    }
  }

  for i in 0..n {
    let to = ans[input_tail_idx[i]];
    let mut ans_str = "";
    if to == undefined { ans_str = "Draw" }
    if to == win { ans_str = "Aoki" }
    if to == lose { ans_str = "Takahashi" }
    println!("{}", ans_str);
  }
}