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

はじめに

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

なぜやるのか

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

目次

A - Distinct Strings (Difficulty 12)
B - Star or Not (Difficulty 62)
C - Calendar Validator (Difficulty 326)
D - Play Train (Difficulty 778)
E - 7 (Difficulty 1678)

前提

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

重複する文字の個数で答えがかわります。
3文字とも同じときは 1
2文字が同じときは 3
ぜんぶ違うときは 6
になります。

fn main() {
  input! {
    mut s: Chars
  }
  s.sort();
  let f1 = s[0] == s[1];
  let f2 = s[1] == s[2];
  if f1 && f2 { // ぜんぶ同じ
    println!("1");
  }
  else if f1 || f2 { // 2文字同じ
    println!("3");
  }
  else { // ぜんぶ違う
    println!("6");
  }
}

B - Star or Not

スターグラフかどうかを判定する問題ですね。
B問題にしてはすこし複雑に感じました。

スターグラフは、エッジが一番多いノードについて、エッジの個数が ノード数 - 1 になります。この特徴を利用します。

実装としては、グラフを構築した後、全ノードについて隣接するエッジの個数を判定しました。

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

  // グラフを構築する
  let mut graph = vec![vec![]; n];
  for (a, b) in ab.into_iter() {
    let a = a - 1;
    let b = b - 1;
    graph[a].push(b);
    graph[b].push(a);
  }

  // サイズが n - 1 のノードがあればスターグラフ
  for v in graph.into_iter() {
    if v.len() == n - 1 {
      println!("Yes");
      return;
    }
  }
  
  // なければスターグラフではない
  println!("No");
}

C - Calendar Validator

バリデーションをする問題ですね。ちょっと複雑です。

方針として、縦方向にチェックした後、横方向にチェックするようにしました。

縦方向のチェックは、すべて等しいかを確認すれば大丈夫です。

横方向のチェックが複雑です。
・数値が 1ずつ増えていること
・2列にまたがっていないこと
を満たしていれば大丈夫です。

例えば、次の2つはどちらも No なので、ちゃんと弾けるか注意です。

1 2 10
8 9 17

13 14 15 16
20 21 22 23

実装です。

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

  let mut ok = true;

  // 縦方向に validation
  for y in 0..n - 1 {
    for x in 0..m {
      ok &= (b[y][x] + 7 == b[y + 1][x]);
    }
  }

  // 0-indexed に直す
  for y in 0..n - 1 {
    for x in 0..m {
      b[y][x] -= 1;
    }
  }

  // 横方向に validation
  // 縦方向の validation ですべて等しいことを判定しているので、b[0] のみを調べればいい
  for x in 0..m - 1 {
    ok &= (b[0][x] + 1 == b[0][x + 1]) && (b[0][x] % 7 < b[0][x + 1] % 7);
  }

  println!("{}", if ok { "Yes" } else { "No" });
}

D - Play Train

計算量の見積もりが必要な問題です。 Vec で愚直にやると時間内に解けません。

まずクエリをひとつずつ見ましょう。

クエリ1 (1 x y)：連結
クエリ2 (2 x y)：分解
クエリ3 (3 x)：出力

3つとも、Vec でやると $O(N)$ かかります。
(クエリ3 は例外ですが、) クエリは Q 回あるので、全体で $O(NQ)$ になります。制約より、時間内には解けません。
ここではクエリ1 とクエリ2 がネックになっているので、連結と分解を高速でできる方法を考えます。クエリ3 には $O(N)$ かけてもいいので、あまり気にしなくても大丈夫です。

連結と分解ができるデータ構造といえば双方向の連結リストですね。これをベースに考えます。連結リストは、各要素が 前 と 次 を持っているので、うまく表現しましょう。

簡単に書きたかったので、構造体を使わずにそれぞれ Vec で表現しました。
こうすることで、連結と分解の両方を $O(1)$ でできるようになり、全体で $O(Q)$ になりました。

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

  let edge = 1001001001; // 前|次 がないときにこの値をいれる

  // prev[i] = 要素i の前
  let mut prev = vec![edge; n];
  // next[i] = 要素i の次
  let mut next = vec![edge; n];

  for _ in 0..q {
    input!{ v: usize }
    if v == 1 { // クエリ1：連結
      input! {
        x: usize,
        y: usize 
      }
      let x = x - 1;
      let y = y - 1;
      next[x] = y;
      prev[y] = x;
    }
    if v == 2 { // クエリ2：分解
      input! {
        x: usize,
        y: usize 
      }
      let x = x - 1;
      let y = y - 1;
      next[x] = edge;
      prev[y] = edge;
    }
    if v == 3 { // クエリ3：出力
      input! { x: usize }
      let mut now = x - 1;
      while prev[now] != edge { // 先頭までいく
        now = prev[now];
      }
      let mut ans = vec![now + 1];
      while next[now] != edge { // 最後尾までいく
        now = next[now];
        ans.push(now + 1);
      }
      print!("{}", ans.len());
      for v in ans.into_iter() {
        print!(" {}", v);
      }
      println!();
    }
  }
}

E - 7

問題文にある通り、原点と (x, y) と (x - 1, y) と (x, y - 1) の四角形を考えます。この四角形が他の四角形と共通部分をもっていなければ問題の条件を満たします。共通部分を持っている場合は、共通部分を取り除くことで条件を満たすようになります。

言い換えると、四角形が共通部分を持っているかどうかを判定する問題になります。
また、個数が大きくなるように工夫して取り除く必要もあります。

四角形の共通部分を考えると難しいので、簡略化を図ります。
すべての四角形は頂点の一つが原点なのでうまく活かしましょう。
考察を進めると、「原点と (x - 1, y) を結ぶ線分」と「原点と (x, y - 1) を結ぶ線分」をみればいいことがわかります。つまり、四角形が三角形になりました。

頂点を共有する三角形同士の共通部分を考えるには、x軸との角度を使った区間を使うとわかりやすいです。
この部分だけを切り取ったら ABC207-C に似た問題になりますでしょうか。この問題を解く人にとっては簡単かとおもいます。

引き続いて考察を進めます。四角形が区間になったので、「重ならないように多くの区間を選ぶ問題」になりました。これは見覚えがありますでしょうか。区間スケジューリング問題ですね。アルゴ式でも取り上げられていて話題になっていました。
簡単に説明すると、区間の後ろでソートして、前から順に貪欲にとるアルゴリズムですね。

今回は区間として角度を使うので、atan2 や分数を使います。私は分数を使いました。分数の場合はゼロ除算に注意です。また、ソートは vec.sort() がそのまま使えないので注意です。

ここまでわかれば実装です。

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

  let mut rad = vec![];
  for (x, y) in xy.into_iter() {
    let x1 = x - 1;
    let y1 = y;

    let x2 = x;
    let y2 = y - 1;

    rad.push( (y1, x1, y2, x2) );
  }

  // ay1/ax1 < by1/bx1 でソート
  // ゼロ除算を避けるために式変形をしている
  rad.sort_by(|(ay1, ax1, ay2, ax2), (by1, bx1, by2, bx2)| {
      if ay1*bx1 == by1*ax1 { (ay2*bx2).cmp(&(by2*ax2)) } else { (ay1*bx1).cmp(&(by1*ax1)) }
    }
  );

  // 区間スケジューリングをする
  let mut now = 0;
  let mut ans = 1;
  for i in 1..n {
    let (_ay1, _ax1, ay2, ax2) = rad[i];
    let (by1, bx1, _by2, _bx2) = rad[now];

    // now_to < i_from
    // by1/bx1 < ay2/ax2
    if by1*ax2 <= bx1*ay2 {
      now = i;
      ans += 1;
    }
  }
  println!("{}", ans);
}

---

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