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

# はじめに

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


# なぜやるのか

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


# 目次

[A - Four Digits](#a---four-digits) (<font color="#808080">Difficulty 5</font>)
[B - Failing Grade](#b---failing-grade) (<font color="#808080">Difficulty 9</font>)
[C - Swiss-System Tournament](#c---swiss-system-tournament) (<font color="#808080">Difficulty 367</font>)
[D - Between Two Arrays](#d---between-two-arrays) (<font color="#008000">Difficulty 865</font>)
[E - Red and Blue Tree](#e---red-and-blue-tree) (<font color="#00C0C0">Difficulty 1491</font>)


# 前提

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

```rust
#[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 - Four Digits

$4 - |N|$ 文字だけ `0` を出力した後、 $N$ を出力します。(全体で複数回出力します。)
$N$ の先頭に `0` を足して、まとめて出力する方法もあります。

```rust
fn main() {
  input! {
    n: Chars
  }
  for _ in 0..4 - n.len() {
    print!("0");
  }
  print!("{}", n.iter().collect::<String>());
}

```


# B - Failing Grade 

$P$ 未満の $a_i$ の個数が答えになりますね。
`iter()` や `into_iter()` などからつなげられる `filter()` と `count()` を使うと簡単に書けます。

```rust
fn main() {
  input! {
    n: usize,
    p: usize,
    a: [usize; n]
  }
  let ans = a.into_iter().filter(|v| v < &p).count();
  println!("{}", ans);
}
```


# C - Swiss-System Tournament 

問題はわかるのですが、実装が複雑ですね。
まず実装はおいておいて、問題にあるじゃんけん大会の流れからみます。

以下の流れになります。

1. 順位で対戦相手を決める (全Nグループ)
1. じゃんけんをする
1. 勝数から順位をつける
1. 次のラウンドに行く

次に、それぞれの処理を考えます。

1. について、`2k` と` 2k + 1` 番目の人が対戦相手になります。
2. について、じゃんけんで何を出すかは入力でもらえますね。勝敗をきめる部分は全通り調べます。
3. について、順位を決めるのは勝数と番号ですね。勝敗の結果と、はじめのインデックスを使ってソートします。
4. について、1~3 を `M`回 くり返します。

3 について、勝数は降順・番号は昇順でソートするので注意が必要です。
私は、 `(勝数 * -1, 番号)` という tuple の Vec をつかいました。こうすることで、 `Vec.sort()` だけで順位をつけられます。

```rust
// c1 と c2 のじゃんけんの結果
// 0 はひきわけ
// 1 は c1 の勝ち
// 2 は c2 の勝ち
fn battle (c1: char, c2: char) -> usize {
  if c1 == c2 { return 0 }
  else if c1 == 'G' && c2 == 'C' { return 1 }
  else if c1 == 'G' && c2 == 'P' { return 2 }
  else if c1 == 'C' && c2 == 'G' { return 2 }
  else if c1 == 'C' && c2 == 'P' { return 1 }
  else if c1 == 'P' && c2 == 'G' { return 1 }
  else if c1 == 'P' && c2 == 'C' { return 2 }

  100 // ここにはこない
}

fn main() {
  input! {
    n: usize,
    m: usize,
    a: [Chars; 2 * n]
  }
  let mut score = vec![(0_i32, 0_i32); 2 * n];
  for i in 0..2 * n {
    score[i].1 = i as i32; // 番号をつける
  }

  for i in 0..m { // 4. M回くりかえす
    for j in 0..n { // 1. Nグループでじゃんけんをする
      let p1 = score[j * 2].1 as usize; // 1. 対戦相手をきめる
      let p2 = score[j * 2 + 1].1 as usize; // 1. 対戦相手をきめる

      let result = battle(a[p1][i], a[p2][i]); // 2. じゃんけんをする
      if result == 1 { score[j * 2].0 -= 1 } // 3. 順位をつける / 勝ち数は -1 をかけるので -1 をしている
      if result == 2 { score[j * 2 + 1].0 -= 1 } // 3. 順位をつける
    }
    score.sort(); // 3. 順位をつける
  }
  for i in 0..n * 2 {
    println!("{}", score[i].1 + 1);
  }
}
```


# D - Between Two Arrays 

DP の問題ですね。


遷移を考えるフェーズ1 と、累積和で高速化するフェーズ2 があります。

まずはフェーズ1 について
$c_{i+1}$ が $M$ の時の個数は、$c_i$ が $M$ 以下の時の個数の総和になりますね。
この遷移をすれば解けるのですが、愚直に実装すると $O(N M^2)$ となり TLE です。

フェーズ2 として高速化をします。
求める値は $0 \sim M$ の和なので、累積和が使えそうです。
累積和の計算量は前処理が $O(M)$ 、区間和の計算が $O(1)$ です。そのため、全体で $O(N M)$ になって時間内に解けました。

続いて実装ですが注意が必要で、$c_i$ の範囲が $a_i \le c_i \le b_i$ なので、$c_i$ から $c_{i+1}$ を計算するときは $a_i \sim b_i$ の範囲の値のみを使うようにします。また、以降の遷移で使うので、累積和は $c_{i+1}$ の値に関係なく $[0, 3000]$ の範囲を計算しています。

```rust
fn main() {
  input! {
    n: usize,
    a: [usize; n],
    b: [usize; n]
  }
  let m = 998244353;
  
  // dp[i] = c の値が i の時の組み合わせ
  let mut dp = vec![1; 3010];

  for i in 0..n {
    let mut next_dp = vec![0; 3010];
    for c in a[i]..=b[i] { // a_i ~ b_i が計算対象になる
      next_dp[c] = dp[c]
    }
    for v in 0..=3000 { // 累積和を求める
      next_dp[v + 1] += next_dp[v];
      next_dp[v + 1] %= m;
    }
    dp = next_dp;
  }
  println!("{}", dp[b[n - 1]]);
}
```

# E - Red and Blue Tree 

グラフとDPが組み合わさった問題です。

まず、それぞれの枝を通る回数を求めます。回数を求めるには、少し難しいですが経路の復元が必要です。
経路を復元するためには、初めての頂点に訪れるたびに直前の枝のインデックスと直前のノードを記録します。目的地についたら、スタート地点まで逆向きにたどっていくことで経路が復元できます。目的地にいくのが $O(N)$ で、スタート地点に戻るのが $O(N)$ なので、全体で $O(N)$ です。
ちなみに [ABC218-F](https://atcoder.jp/contests/abc218/tasks/abc218_f) でも似た処理が出ました。
これを $M$ 回やるので、この部分の計算量は $O(NM)$ になります。

次はDPです。
`dp[i] = 和が i になるときの組み合わせ` という dpテーブルで解けます。これはよくある DP で、 [ABC216-F](https://atcoder.jp/contests/abc216/tasks/abc216_f) でも出ました。
ちなみに、`dp[i][j] = 枝 i まで見て和が j になるときの組み合わせ` とするとメモリ制限超過になるので注意です。

DPテーブルはできましたが、 `k` は負の値もとるので困りますね。こういうときは、 `dp[k]` を「和が `0` のとき」として扱います。
`dp[0]` が 和が「`-k` のとき」を示して、 `dp[2*k]` が「和が `k` のとき」を示すようにします。
こうすると、配列のインデックスが負にならずに済みます。
また、`K` の代わりに、取り得る最大値 (コードで言う `use_count.iter().sum();` ) を使うことで、少しばかりのメモリ節約と高速化を狙っています。

DP部分の計算量は $O(N |K|)$ になります。
全体の計算量は $O(NM + N |K|)$ となりました。

```rust
fn main() {
  input! {
    n: usize,
    m: usize,
    k: i64,
    a: [usize; m],
    uv: [(usize, usize); n - 1]
  }

  let mut graph = vec![vec![]; n];
  for i in 0..n - 1 {
    let (u, v) = uv[i];
    let (u, v) = (u - 1, v - 1);
    graph[u].push((v, i));
    graph[v].push((u, i));
  }

  let inf = 1001001001;
  let mut use_count = vec![0; n - 1]; // 枝を使った回数

  // from から to までの最短経路で通る枝を記録する
  let mut bfs_shortest_path_index = |from: usize, to: usize| {
    // prev[i] = (頂点, index)
    let mut prev = vec![(inf, inf); n];
    let mut dist = vec![inf; n];
    dist[from] = 0;

    let mut que = VecDeque::new();
    que.push_back(from);

    // BFS で最短経路を求める
    while let Some(now) = que.pop_front() {
      let depth = dist[now];
      if now == to { break }
      for &(next, idx) in graph[now].iter() {
        if depth + 1 < dist[next] {
          dist[next] = depth + 1;
          prev[next] = (now, idx);
          que.push_back(next);
        }
      }
    }

    // 使った枝を記録する
    let mut now = to;
    while now != from {
      use_count[prev[now].1] += 1;
      now = prev[now].0;
    }
  };

  for i in 1..m {
    bfs_shortest_path_index(a[i - 1] - 1, a[i] - 1);
  }

  let m = 998244353_usize;
  use_count.sort(); // ソートして小さい順に処理する
  let sum: usize = use_count.iter().sum();
  let mut sumv = 0;

  if max(k, -k) as usize > sum {
    println!("{}", 0);
    return;
  }

  // dp[i] = 和が i になるときの組み合わせ
  let mut dp = vec![0_usize; sum * 2 + 1_usize];
  dp[sum] = 1; // dp[sum] が 和が 1 の時の組み合わせを示す
  for i in 0..n - 1 {
    let mut next_dp = vec![0_usize; sum * 2 + 1_usize];
    sumv += use_count[i]; // これまでの累積和
    for j in sum - sumv..=sum + sumv { // ソートして小さい順に処理することで、dpテーブルの更新範囲を小さくする
      if j + use_count[i] <= sum * 2 {
        next_dp[j] += dp[j + use_count[i]];
        next_dp[j] %= m;
      }
      if j >= use_count[i] {
        next_dp[j] += dp[j - use_count[i]];
        next_dp[j] %= m;
      }
    }
    dp = next_dp;
  }
  println!("{}", dp[(sum as i64 + k) as usize]);
}
```



----

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