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

はじめに

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

なぜやるのか

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

目次

A - Lexicographic Order (Difficulty 21)
B - AtCoder Quiz (Difficulty 22)
C - Inverse of Permutation (Difficulty 49)
D - Cutting Woods (Difficulty 802)
E - Sorting Queries (Difficulty 986)
F - Make Pair (Difficulty 1954)

前提

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

文字の辞書順での大小を比較する問題ですね。
Rust は String 型の比較演算子で、問題通りの大小比較ができます。

fn main() {
  input! {
    s: String,
    t: String
   }
   let ans = if s < t { "Yes" } else { "No" };
   println!("{}", ans);
}

B - AtCoder Quiz

特定の文字列と入力文字列を比較して、入力文字列にないものを答える問題ですね。

私は特定の文字列を配列で持って、1つずつ条件を満たすかどうかを判定しました。
他にも HashSet を使う方法や HashMap を使う方法など、いろいろな解き方があります。

fn main() {
  input! {
    s: [String; 3]
  }
  let contests = vec!["ABC", "ARC", "AGC", "AHC"];
  for c in contests {
    if !s.contains(&c.to_string()) {
      println!("{}", c);
    }
  }
}

C - Inverse of Permutation

順列 $P$ の値とインデックスを入れ替える問題ですね。
問題文がすこし難解に感じてしまいました。

enumerate() をするとインデックスをつけられます。
入れ替えた後のインデックス順に出力するので、インデックスでソートしましょう。

enumerate() をすると (index, value) の順になるのですが、今回は入れ替えるので (val, idx) としています(変数名の付け方の話です)。
また、 enumerate() で作られるインデックスは 0-index ですが、答えは 1-index なので注意です。

fn main() {
  input! {
    n: usize,
    mut p: [usize; n]
  }
  let mut ans = p
    .into_iter()
    .enumerate()
    .collect_vec();

  ans.sort_unstable_by_key(|(_val, idx)| *idx);
  for (val, _idx) in ans {
    println!("{}", val + 1);
  }
}

D - Cutting Woods

クエリの問題ですね。ひとつずつ見ていきましょう。

クエリ1 : x に印をつける
クエリ2 : x から前と後ろにある印の距離を計算する

クエリ1 はできそうですが、 クエリ2 は少し手間がかかりそうですね。というのも、愚直にやるとクエリひとつあたり $O(Q)$ かかって、全体では $O(Q^2)$ になります。制約から、これでは実行時間に間に合いません。

ここでは二分探索を使いましょう。二分探索では、単調増加の列に対して、「ある値以上の最小の数」を高速に求められます。「ある値以上の最小の数」は今回の問題で言う、クエリ2 の「後ろの位置」にあたります。さらに、そのひとつ前の要素が、クエリ2 の「前の位置」になります。

また、クエリ2 に対応するため、クエリ1 でつける印は昇順になるようにします。
これを効率的にやるには挿入ソートの挙動になるのですが、例えば挿入と二分探索ができる Vec でやろうとすると、Vec の要素数を $N$ として、挿入に $O(N)$ かかります。
この場合も全体では $O(Q^2)$ になって、制約から実行時間に間に合いません。

さて、昇順になるような挿入と探高速な探索をするにはどうすれば良いでしょうか。
これらを達成するには、平衡二分探索木を使います。Rust では BTreeSet になります。

クエリ1 は insert() 関数を使います。 $O(log N)$ で、昇順でソートされます。
最先端の 0 と最後尾の l を先に挿入しておくことに注意です。

クエリ2 はトリッキーですが、 range() 関数を使います。戻り値が double-ended iterator です。
・range(..x).next_back().unwrap() で 「 x 未満の最大の数」
・range(x..).next().unwrap() で 「 x 以上の最小の数」
を取得できます。

さて、これで クエリ1 とクエリ2 をどちらも $O(\log Q)$ で処理できるようになりました。
全体で $O(Q\ \log Q)$ になって、時間内に解けました。

fn main() {
  input! {
    l: usize,
    q: usize,
    query: [(usize, usize); q]
  }

  let mut st: BTreeSet<usize> = BTreeSet::new();
  st.insert(0);
  st.insert(l);
  for (c, x) in query {
    if c == 1 {
      st.insert(x);
    }
    else {
      let lower = st.range(..x).next_back().unwrap();
      let upper = st.range(x..).next().unwrap();
      println!("{}", upper - lower);
    }
  }
}

E - Sorting Queries

またクエリの問題ですね。またひとつずつ見ていきましょう。

クエリ1 : 列の後ろに追加する
クエリ2 : 列の先頭を pop する
クエリ3 : 列をソートする

クエリの 1 と 2 を見ると deque が使えそうです。クエリの 2 と 3 を見ると heap が使えそうです。ですが、クエリの 1 と 2 と 3 を一緒に扱えるデータ構造となると思いつかないです。というのも、クエリの 1 と 3 があるために、ソート済みと未ソートの状態が混在します。
Vec を使えばすべての操作ができるかもしれませんが、先頭の pop に $O(Q)$、 ソートするのに $O(Q\ \log Q)$ かかって、全体で $O(Q^2\ \log Q)$ かかります。今回の制約からこれは実行時間に間に合いませんね。

さて、できないのならわけてしまいましょう。
・クエリの 1 と 2 を deque
・クエリの 2 と 3 を heap
で扱います。

こうするとすべてのクエリを高速に扱えます。

クエリ1 : deque の後ろに追加する -> $O(1)$
クエリ2 : heap の先頭を pop する -> $O(log Q)$
クエリ3 : deque の要素をすべて heap に挿入する -> $O(\log Q)$ が最大で $Q$ 回

これで、全体で $(Q \log Q)$ の計算量になって時間内に解けました。

ちなみに、Rust の deque は VecDeque、 heap は BinaryHeap を使います。
BinaryHeap はデフォルトでは最大ヒープなので、 Reverse を使って最小ヒープにしています。

use std::collections::BinaryHeap;
use std::cmp::Reverse;

fn main() {
  input! {
    q: usize
  }

  let mut pool = VecDeque::new();
  let mut heap = BinaryHeap::new();
  for _ in 0..q {
    input!(n: usize);
    if n == 1 {
      input!(x: usize);
      pool.push_back(x);
    }
    if n == 2 {
      if let Some(Reverse(top)) = heap.pop() {
        println!("{}", top);
      }
      else if let Some(top) = pool.pop_front() {
        println!("{}", top);
      }
    }
    if n == 3 {
      while let Some(v) = pool.pop_front() {
        heap.push(Reverse(v));
      }
    }
  }
}

F - Make Pair

区間DP の問題です。難しいです。

今回は解説ACをして、区間DPをどうやって思いつくかわかりませんでした。
DPなので、思いつくには経験が必要なんじゃないかとは思います。
考え方とかはあまりかけておらず、実装したことそのままの説明になっています。。。

さて、DPテーブルですが、 dp[left][right] = 区間 [left, right] の組み合わせ とします。(ちなみに [left, right] は閉区間の記法です。)

まず left と right を区間の幅が小さいものから処理します。つまり、作るペアの数が小さいものから処理します。区間の組み合わせを計算するときに、より幅の狭い区間の組み合わせが必要なためですね。
次に、 left とペアになる相手を left + 1 から right までひとつずつ決めて組み合わせを計算します。この時 left の相手を p とすると、 [left, p] の幅が奇数の状況は起こり得ないので飛ばせます。

こうすると、できるペアは

1. left と p の 1 ペアで1通り
2. [left + 1, p - 1] の (p - left - 1) / 2 ペアで dp[left + 1][p - 1] 通り
3. [p + 1, right] の (right - p) / 2 ペアで dp[p + 1][right] 通り
   あります。

dp[l][r] のインデックスは l < r になるようにしているのですが、 dp[left + 1][p - 1] と dp[p + 1][right] はこの大小関係が崩れるタイミングがあります。その時に値がおかしくならないように注意が必要です。

さて、 1 のペアを作るのは 2 のペアをすべて作った後になります。ここは順番が決まっています。
ですが、 1,2 と 3 には順番がありません。つまり、ここはペアを作る順列が複数通りあります。
1,2,3 の合計ペア数を n 、 3 のペア数を r とすると、組み合わせは ${}_n C_r$ 通りですね。

これらをかけると dp[left + 1][p - 1] * dp[p + 1][right] * nCrになります。
これで、 [left, right] の範囲で left の相手を p にしてペアを作るときの組み合わせが何通りあるかわかります。
さらに、 p を left からright まで動かしながら計算すれば [left, right] の範囲でペアを作る組み合わせが何通りかわかります。

また、計算量を考えると [left, right] の範囲が $O(N)$ で、left と right を全通り計算するので $O(N^2)$ になります。
これで、全体で $O(N^3)$ となり、時間内に解けました。

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

  let m: i64 = 998244353;
  let sz = n * 2;
  // comb[n][k] = n から ｋ とる組み合わせ () n C k ) の値
  let mut comb = vec![vec![0_i64; sz + 1]; sz + 1];
  comb[0][0] = 1;
  for i in 1..sz {
    comb[i][0] = 1;
    comb[i][i] = 1;
    for j in 0..i - 1 {
      comb[i][j + 1] = comb[i - 1][j] + comb[i - 1][j + 1];
      comb[i][j + 1] %= m;
    }
  }

  let mut ok = vec![vec![false; sz]; sz];
  for (a, b) in ab {
    let (a, b) = (a - 1, b - 1);
    ok[a][b] = true;
  }

  // dp[left][right] = 区間 [left, right] の組み合わせ
  let mut dp = vec![vec![0_i64; sz]; sz];
  for pair in 1..=n { // つくるペアの個数
    for left in 0..sz { // 左端のインデックス
      let right = left + pair * 2 - 1; // 右端のインデックス
      if right >= sz { break; } // 範囲外で break

      for i in 0..pair {
        let p = left + i * 2 + 1; // left とペアにするインデックス
        if !ok[left][p] { continue; } // ペアになれなければ continue

        let l = if left + 1 < p - 1 { dp[left + 1][p - 1] } else { 1 }; // p の左側のペア数
        let r = if p + 1 < right { dp[p + 1][right] } else { 1 }; // p の右側のペア数
        let perm = l * r % m;
        dp[left][right] += perm * comb[pair][i + 1]; // [left][iペア][p][pair - i - 1 ペア]
        dp[left][right] %= m;
      }
    }
  }
  println!("{}", dp[0][sz - 1]);
}

---

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