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

はじめに

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

なぜやるのか

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

目次

A - Weather Forecast (Difficulty 7)
B - qwerty (Difficulty 14)
C - Shapes (Difficulty 1012)
D - Rectangles (Difficulty 715)
E - Destruction (Difficulty 1004)
F - Blocked Roads (Difficulty 1753)

前提

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

入力文字列 S の N 文字目を判定する問題ですね。

S を Chars などの配列でもらって N 番目の要素を判定しましょう。
0-index になおすことに注意です。

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

  let ans = if s[n - 1] == 'o' { "Yes" } else { "No" };
  println!("{}", ans);
}

B - qwerty

数字を文字に変換する問題ですね。

文字 (char) は内部的には数字であることを利用して解きます。usize の配列として受け取った後、1 が a に対応するように変換します。コードでは (p - 1) + 'a' as usize の部分です。

fn main() {
  input! {
    p: [usize; 26]
  }
  let ans = p
    .into_iter()
    .map(|p| (p - 1) + 'a' as usize)
    .map(|p| p as u8 as char)
    .join("");
  println!("{}", ans);
}

C - Shapes

図形を回転しながらマッチングする問題です。

まずは、回転とマッチングを分けて考えましょう。初めに回転です。

回転は、 (x, y) が (n - y, x) に移動しますが、これは回転行列を使って計算します。
点(x, y) を (n/2, n/2) を中心に90度回転させる計算で、以下の式になります。
(アフィン変換を使えば、平行移動も行列同士の乗算で表現できます。興味ありましたら調べてみてください。)

\begin{pmatrix}
x' \\
y' \\
\end{pmatrix}
=
\begin{pmatrix}
cos \frac{\pi}{2} & -sin \frac{\pi}{2} \\
sin \frac{\pi}{2} & cos \frac{\pi}{2} \\
\end{pmatrix}
\begin{pmatrix}
x - n/2 \\
y - n/2 \\
\end{pmatrix}
+
\begin{pmatrix}
n/2 \\
n/2 \\
\end{pmatrix}
=
\begin{pmatrix}
0 & -1 \\
1 & 0 \\
\end{pmatrix}
\begin{pmatrix}
x - n/2 \\
y - n/2 \\
\end{pmatrix}
+
\begin{pmatrix}
n/2 \\
n/2 \\
\end{pmatrix}
=
\begin{pmatrix}
n - y \\
x \\
\end{pmatrix}

...まあ、ここまでかっちり考えなくても
(0, 0) -> (0, n)
(0, n) -> (n, n)
(n, n) -> (n, 0)
(n, 0) -> (0, 0)
となる変換式を考えれば導出できます。
y = x の一次関数を y = n - x の一次関数に変換することをイメージしてもいいかもしれません。
私はいつも、行列を使わずにこれらの方法でやっています。

次に、マッチングについてです。
マッチングは、図形をできるだけ左上に詰めて判定します。こうすることで、図形が全体のどこにいるかに関係なくマッチングができます (詰めずに判定すると判定回数が増えて TLE になります)。
この詰める量は、全要素のうち最小の x と y になりますね。コードで言う min_ty, min_tx, min_sy, min_sx です。

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

  // Vec<string> を Vec<Vec<char>> にする
  let mut s = s.into_iter().map(|s| s.chars().collect_vec()).collect_vec();
  let t = t.into_iter().map(|t| t.chars().collect_vec()).collect_vec();

  let mut tp = Vec::<(usize, usize)>::new();
  let mut min_ty: usize = 1001001001;
  let mut min_tx: usize = 1001001001;

  for y in 0..n {
    for x in 0..n {
      if t[y][x] == '#' {
        tp.push((y, x));
        min_tx = min(min_tx, x);
        min_ty = min(min_ty, y);
      }
    }
  }

  let tp = tp
    .into_iter()
    .map(|(y, x)| (y - min_ty, x - min_tx))
    .collect_vec();

  for _ in 0..4 {
    let mut sp = Vec::<(usize, usize)>::new();
    let mut min_sy: usize = 1001001001;
    let mut min_sx: usize = 1001001001;
    for y in 0..n {
      for x in 0..n {
        if s[y][x] == '#' {
          sp.push((y, x));
          min_sx = min(min_sx, x);
          min_sy = min(min_sy, y);
        }
      }
    }

    let sp = sp
      .into_iter()
      .map(|(y, x)| (y - min_sy, x - min_sx))
      .collect_vec();

    if sp == tp {
      println!("Yes"); // 一致した
      return;
    }

    let mut rot = vec![vec!['.'; n]; n]; // 90度回転
    for y in 0..n {
      for x in 0..n {
        rot[y][x] = s[x][n - y - 1];
      }
    }
    s = rot;
  }
  println!("No");
}

D - Rectangles

x軸とy軸に平行な辺を持つ正方形を見つける問題です。

全探索をすると、全部で ${}_N C_4$ 通りで $O(N^4)$ となって時間内には解けません。
私は、2点を選んで、その2点を結ぶ線分を対角線とする正方形をみつける方法で解きました。対角線となる2点が決まれば、のこりの2点の座標はわかりますね。
というのも (x1, y1) と (x2, y2) を選ぶと、残りの2点は (x1, y2) と (x2, y1) になります(図に書いてみると分かりやすいかと思います)。

2点を選ぶ方法は ${}_N C_2$ 通りで $O(N^2)$ なので、残りの2頂点を高速で見つけられれば時間内に解くことができます。
この2点を高速で見つけるために、平衡二分探索木を使いました。Rust では HashSet です。座標をタプルにしたものを HashSet にいれれば、高速で検索できます。コードでいう以下の部分です。

if let Some(_) = st.get(&(xi, yj))

この検索は $O(\log N)$ でできるので、全体では $O(N^2 \log N)$ の計算量で解けました。

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

  xy.sort_unstable();

  let mut st = HashSet::<(usize, usize)>::new();
  for v in xy.iter() {
    st.insert(*v);
  }

  let mut ans = 0;
  for i in 0..n {
    let (xi, yi) = xy[i];
    for j in i + 1..n {
      let (xj, yj) = xy[j];
      
      // (x_i, y_i) と (x_j, y_j) の2点を選ぶ
      if xi < xj && yi < yj {
        if let Some(_) = st.get(&(xi, yj)) { // (x_i, y_j) がみつかる
          if let Some(_) = st.get(&(xj, yi)) { // (x_j, y_i) がみつかる
            ans += 1;
          }
        }
      }
    }
  }
  println!("{}", ans);
}

E - Destruction

グラフの問題ですね。

グラフの辺をつないだり切ったりする点から、 最小全域木や UnionFind木 が連想されます。
ABC214-D を、UnionFind木の問題だと気付きやすくした問題だと感じました。

ただ今回はそれとは違って、枝を切ると枝の点数がもらえます。最大化するには、正の点数の枝を切れるだけ切って、負の点数の枝はできるだけ残せば良さそうです。
もうちょっと具体的に、どんな状態だと点数が最大になるかを考えてみましょう。結論から言うと、より点数の低い枝だけで連結になっている状態ですね。点数の低い順に枝を繋げていって、グラフを連結にする処理が思い浮かびます。つまり、枝を切るのではなく繋いでいきます。ここで UnionFind木がでてきます。

UnionFind木を思いつけば、あとはコストの少ない順に繋いでいけば答えがわかります。連結になった後でも、コストが負の枝はつないだ方が全体のスコアが高くなることに注意です。

実装では、 petgraph::unionfind::UnionFind を使いました。

use petgraph::unionfind::UnionFind;

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

  let mut cab = abc
    .into_iter()
    .map(|(a, b, c)| (c, a - 1, b - 1))
    .collect_vec();

  cab.sort_unstable();

  let mut ans: i64 = cab
    .iter()
    .map(|(c, _, _)| c)
    .sum();

  let mut uf = UnionFind::<usize>::new(n);
  for (cost, from, to) in cab.into_iter() {
    if !uf.equiv(from, to) {
      uf.union(from, to);
      ans -= cost;
    } else if cost < 0 {
      ans -= cost;
    }
  }
  println!("{}", ans);
}

F - Blocked Roads

グラフ上の距離を答える問題です。

すべての辺を使えるときは BFS または DFS を1回すれば答えがわかりますね。この問題で難しいのは、特定の辺が使えないときに、これまでの計算結果をうまく使って距離をだせない、という点です。そう。うまく使っても距離をだせません。だせないので、もういちど BFS か DFS をする必要があります。
ですが、いちどの探索に $O(M)$ かかって、辺は $M$ 個あるので、全体の計算量が $O(M^2)$ になります。これでは時間内に解けません。困りました。

ここでひとつ大事なのは、特定の辺を切ったときに、「その辺が頂点 1 から 頂点 N の最短経路上にある」場合にのみ、「最短経路が変わる」という点です。(ド・モルガンの法則で) 言い換えると、「特定の辺が頂点 1 から 頂点 N の最短経路上にない」場合は、「最短経路は変わりません」。
つまり、最短経路が変わるときは再計算して、最短経路が変わらないときは計算をしないという方針が取れます。
こうすると嬉しいことがあって、最短経路の距離は最長でも $N$ になります。つまり、再計算の回数が最大でも $N$ 回になります。

最大で $N$ 回は $O(M)$ で探索して、そのほかの場合は $O(1)$ で答えが出せるので、この問題は全体で $O(NM)$ になりました。

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

  let st = st
    .into_iter()
    .map(|(s, t)| (s - 1, t - 1))
    .enumerate()
    .collect_vec();
  
  let mut graph = vec![vec![]; n];
  for &(idx, (s, t)) in st.iter() {
    graph[s].push((t, idx));
  }

  let inf = 1001001001;
  let bfs = |invalid_idx: usize| {
    // vec[(depth, from_edge_idx)]
    let mut dist = vec![(inf, 0); n];
    dist[0].0 = 0;

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

    while let Some(now) = que.pop_front() {
      let depth = dist[now].0;
      for &(to, idx) in graph[now].iter() {
        if idx != invalid_idx && depth + 1 < dist[to].0 {
          dist[to] = (depth + 1, idx);
          que.push_back(to);
        }
      }
    }
    dist
  };

  let dist = bfs(m);
  if dist[n - 1].0 == inf { // たどり着けない場合は -1 を出力して終了
    for _ in 0..m {
      println!("-1");
    }
    return;
  }

  let mut shortest_path = vec![false; m];
  let mut now = n - 1;
  while now != 0 { // 最短経路に含まれる枝を記録する
    let edge_idx = dist[now].1;
    shortest_path[edge_idx] = true;
    now = (st[edge_idx].1).0;
  }

  for i in 0..m {
    let mut ans = dist[n - 1].0;
    if shortest_path[i] {
      let d = bfs(i);

      ans = if d[n - 1].0 == inf {
        -1
      } else {
        d[n - 1].0
      };
    }
    println!("{}", ans);
  }
}

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