Rust入門日記 - 競技プログラミング編
ooharabucyou
の続き
Rust の基本的な動きを確認した上で、実際にAtCoder (ABC) の問題を取り扱う。
流れとしては
- 標準入出力
- ファイル入出力
をはじめにやる。競技プログラミングの場合は、便利なパッケージを利用するため、このへんは深く考えなくても良さそうだが、競技プログラミングに問わず利用できそうな話なので、一度パッケージを使わずに、Rust標準の機能を利用して行う。
その後、proconio という、競技プログラミング用のパッケージを利用して、問題をいくつか回答してみる。
ooharabucyou標準入力
問題例
例によって、以下の問題を解くことを考える。
高橋くんと青木くんがゲームを行います。
はじめ、高橋くんは A個、青木くんは B個のアメを持っています。
C=0 ならば高橋くんが先手、C=1 ならば青木くんが先手で、高橋くんと青木くんは以下の操作を交互に繰り返します。
- 自分の持っているアメを 1個食べる。
先に操作を行えなくなった者の負けです。どちらが勝つでしょうか？
入力は、標準入力によって、以下の形式で行われるとのことです。
A B C
出力は、 高橋くんが勝つならば Takahashi を、青木くんが勝つならば Aoki を出力することのことです。
実装を考える
まず、以下のことが行える必要がある。
- 標準入力から文字列を1行入力する (競技プログラミングの上ではエラーはないものとする。)
- 入力した文字列を " " で、分割する
- 分割した文字列を、数値として解釈する
また、この問題については、以下のように解くことができる。
- 文意通り高橋くんと青木くんのアメカウンタを作って、交互に消費していき0未満になったら負けとする
- 先攻が高橋くんだったとき、青木くんが高橋くん以上のアメを持っていたら青木くんの勝ち。そうでなければ、高橋くんの勝ち。先攻が青木くんだった場合は、高橋くんが青木くん以上のアメを持っていたら高橋くんの勝ち。そうでなければ青木くんの勝ち。
どう考えても、繰り返し等を行わなくても良い、後者のほうが効率が良いので、後者で解くことにする。
標準入力のやり方
ということで、はじめての標準入力。std::io::stdin().read_line(&mut 変数) は変更可能な String を用意しておいて、そこに1行取り込んで入れ込むというもの。
とはいえ、入力のエラーも想定する必要があるので、帰り値は、std::io::Result列挙体となり、enum Result<T, E> と同じく、Ok か、 Err の状態になる。Ok な場合入力した bytes が入っていて、Err の場合、std::io::Error の型となる。
fn main() {
let mut s: String = String::new();
let result: std::io::Result<usize> = std::io::stdin().read_line(&mut s);
match result {
Ok(_) => {
println!("{}", s);
}
Err(err) => {
println!("{}", err);
}
}
}
パイプで標準入力のテスト
$ cargo build
$ echo "hoge" | ./target/debug/hello_cargo
hoge
ただし、競技プログラミングでは標準入力のエラーは想定しなくてよいので、以下のようにする。
std::io::stdin().read_line(&mut s).ok();
ok() は、Result::Ok の値のみを取り出す Result に実装されたメソッド。ok() はなくても一応動作するが、警告は出るので書いておいたほうが良いだろう。
これで、s には、入力した文字列が入るので、これをスペースで分割して数値に変換した Vec を作りたい。
ということで、String にできる操作を駆使していく。
let v:Vec<u8> = s.trim().split_whitespace().map(|e| e.parse().ok().unwrap()).collect();
一見よくわからないが分解してみていく。
trim() は他の言語でもよくありがちな、文の初めや終わりの余計な文字を除去したものを返すメソッド。
その後、split_whitespace() により、スペースにより文字列を分割して、SplitWhitespace という型を返す。
これは、繰り返し処理を行う機能がついているので、それに対して利用できる、map() を利用する。(map は Ruby などにあるのと同じ。)
分割された文字列 (&str) を parse() して、数値として解釈する。これは Result<T, E> を返す。これも競技プログラミング的には例外はないので、ok() な値のみを取り出す。
これは、Option 型で、Some のみを取り出すために、unwrap() を利用する。
map() の結果は collect() でまとめで、Vec<?> として v に入れる。
という形で、値の入力できた。
あとは、考えた解答法に応じてプログラムを書くだけ。
fn main() {
let mut s: String = String::new();
std::io::stdin().read_line(&mut s).ok();
let v: Vec<u8> = s.trim().split_whitespace().map(|e| e.parse().ok().unwrap()).collect();
if v[2] == 0 {
if v[0] <= v[1] {
println!("Aoki");
} else {
println!("Takahashi");
}
} else {
if v[0] >= v[1] {
println!("Takahashi");
} else {
println!("Aoki");
}
}
}
$ cargo build
$ echo "2 2 0" | ./target/debug/hello_cargo
Aoki
実際には、Cの値については同じだったときだけ配慮すればいいので、以下のように書いてもよさそうだ。
fn main() {
let mut s: String = String::new();
std::io::stdin().read_line(&mut s).ok();
let v: Vec<u8> = s.trim().split_whitespace().map(|e| e.parse().ok().unwrap()).collect();
if v[0] > v[1] {
println!("Takahashi");
} else if v[0] < v[1] {
println!("Aoki");
} else if v[2] == 0 {
println!("Aoki");
} else if v[2] == 1 {
println!("Takahashi");
}
}
$ cargo build
$ echo "2 2 0" | ./target/debug/hello_cargo
Aoki
ファイルの入出力
AtCoder競技プログラミングにおいては、ファイルの入出力についてやる必要はないが、Rust そのものの勉強も兼ねているのでやることにする。
ファイルの入力 (全部読む)
ファイルの入力については、標準で読み込まれていないモジュールを利用する。use を利用して、利用する機能を呼び出すことができる。Javaなどと同様、File を使う場合に std::fs::File などとフルで書いても良い。
std::io::Read については、入力に関するトレイトで、ファイルに read_to_string() などの便利なメソッドを用意してくれる。
use std::fs::File;
use std::io::Read;
fn main() {
let mut f = File::open("./sample.txt").expect("ファイルが見つかりませんでした");
let mut contents = String::new();
f.read_to_string(&mut contents).expect("ファイルの入力に失敗しました");
println!("{}", contents);
}
File構造体は、入力処理時に中身をいじることになるので、mut で定義する必要がある。
標準入力と同じく、可変変数の文字列を用意して、そこにファイルの中身を書き出すという仕組みになっている。
今回、expect() を使っているが、これは、 Result が Err となった場合に、指定したエラーメッセージを出力し、panic! を起こしてプログラムを止めてくれるというもの。便利やん。
read_to_string() は、ファイルから中身を全部読んで文字列として読み取ってくれる大変便利なメソッド。PHPの file_get_contents() ばりに便利な機能。巨大はファイルの読み取りは当然ながらヒープ領域をたくさん使うことになりそうなので、別の方法を使ったほうが良さそうだ。
ファイルの入力 (1行づつ読む)
エラー処理をあまりやらない前提で以下。構造体BufReader と、BufReadトレイトを使う。同じパッケージ内にある、2つのモジュールを使いたい場合は、 use std::io::{BufRead, BufReader}; のように指定できる。
lines() の結果を for で取り出すことにより、1行づつ取得できる。
以下の場合、line には、 std::io::Result<String> が入っている状態になり、ファイルが最後まで読み取れるまで繰り返す。
unwrap() は Result 列挙体 (の系列) に行える処理の一つで、エラーが発生していない場合はそのまま Ok の中身を返して、エラーが起きている場合は panic! を起こします。
use std::fs::File;
use std::io::{BufRead, BufReader};
fn main() {
let f = File::open("./sample.txt").expect("ファイルが見つかりませんでした");
for line in BufReader::new(f).lines()
{
println!("{}", line.unwrap())
}
}
最初の一行だけ読みたいという場合は、read_line() を利用する。
すべての内容を読むのではなく、指定の行数を読んでおく場合などはこちら。
use std::fs::File;
use std::io::{BufRead, BufReader};
fn main() {
let f = File::open("./sample.txt").expect("ファイルが見つかりませんでした");
let mut buff = BufReader::new(f);
let mut line = String::new();
buff.read_line(&mut line).expect("ファイルの読み取りに失敗しました");
println!("{}", line);
}
ファイルの入力 (その他)
あとは、バイナリをそのまま読み出す方法やら、バイナリをバッファリングしながら読む方法があるが、しばらく使う予定はないので、必要なときに読むことにする。
参考記事
ooharabucyouファイルの出力 (1行書き出す)
std::io::Write にある、writeln! マクロを使うことで、ファイルに対して文字列を書き出すことができる。
println! と同様、writeln!(file, "v={}", v); のように変数を埋め込むこともできる。
use std::fs::File;
use std::io::Write;
fn main() {
let mut file = File::create("./sample.txt").expect("ファイル作成に失敗");
writeln!(file, "hello").expect("書き出しエラー");
}
ファイル出力 (Append)
他のプログラミング言語では、ファイルオープンモードなどのフラグを渡すことによって、ファイルに追記を行ったりできる。上記の方法だと、ファイルがまっさらになってしまうので、追記する形にしたい。
以下のように、OpenOptions を使って、書き出すと1行つづ出力される。
use std::fs::OpenOptions;
use std::io::Write;
fn main() {
let mut file = OpenOptions::new()
.append(true)
.create(true)
.open("./sample.txt").expect("ファイル作成エラー");
writeln!(file, "hello").expect("書き出しエラー");
}
上記のように、シンプルなものであれば問題ないが、ある程度多くの量のデータを書き出すとなると、パフォーマンス的な問題になってくる。
例えば、以下のように 100万行のデータを書き出すプログラムを書くとして、実行時間を測定すると、0.43秒となる。
for _ in 0..1_000_000 {
writeln!(file, "hello").expect("書き出しエラー");
}
./target/debug/hello_cargo 0.43s user 2.17s system 98% cpu 2.629 total
以下のように、バッファリングを通して書き出しを行うようにプログラムを変える。
use std::fs::OpenOptions;
use std::io::{BufWriter, Write};
fn main() {
let file = OpenOptions::new()
.append(true)
.create(true)
.open("./sample.txt").expect("ファイル作成エラー");
let mut buf_writer = BufWriter::new(file);
for _ in 0..1_000_000 {
writeln!(buf_writer, "hello").expect("書き出しエラー");
}
}
これにより、4倍近くパフォーマンスが改善される。
Rust は、標準のファイルをそのまま利用する場合だと、バッファリングを行ってくれないので、特別な事情がない場合、読み取りは BufReader、書き出しは BufWriter を利用して書き出すことを心がけたほうが良さそうではある。
./target/debug/hello_cargo 0.11s user 0.01s system 98% cpu 0.115 total
なお、ライフサイクルが終了したタイミングで、ファイルを自動的に閉じてくれるので、close は原則的に行う必要はないようだ。便利である。
ooharabucyouということで、仮に proconio を使わない前提でファイル入出力についても考えてみた。
次回は、proconio を利用して実際に問題を解いてみる。
ooharabucyouproconio を使う (1)
AtCoder では、入力を便利にするための proconio というライブラリを使うことができる。
Cargo.toml の記述
Cargo.toml には、必要なライブラリを記述することができる。Gemfile みたいなものだろうか。
に習って、0.3.6 を使う設定としてみる。
[dependencies]
proconio = "=0.3.6"
利用
入力の形式が
A B
で、A, Bがともに数値の場合は以下のように記述するだけで入力ができる。
input! マクロでは、変数の定義だの、std::io を利用した入力処理などに置き換えてくれる。
use proconio::input;
fn main() {
input! {
a: i32,
b: i32
}
println!("{}", a * b)
}
試しに実行する。
$ echo "2 2" | cargo run
4
proconioを使う (2)
もう少し複雑な入力形式を考えてみる
N1 N2
i1_1 i1_2 i1_3 ... i1_N2
i2_1 i2_2 i2_3 ... i2_N2
...
iN1_1 iN1_2 iN1_3 ... iN1_N2
のような形式で、N1には行数、N2には列数が指定された、行列I (それぞれの要素は数値) を入力したいとする。
サンプルとして以下のデータを用意する。
4 5
1 2 3 4 5
6 7 8 9 10
11 12 13 14 15
16 17 18 19 20
これを以下のように利用できる。
use proconio::input;
fn main() {
input! {
n1: i32,
n2: i32,
i: [[i32; n2]; n1]
}
println!("{:?}", i)
}
i は、Vec<Vec<i32>> 型となり、実行すると以下のようになる。
$ cat hoge.txt| cargo run
[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10], [11, 12, 13, 14, 15], [16, 17, 18, 19, 20]]
proconio を使う (3)
数値だけではなく文字列を入力することもできる。
以下のような形式のデータを読み込むとする。
n は数値、Sは文字列とする。
3
hello
world
atcoder
これも、単純に String と書いておけば、文字列で入力できる。
下記のコードでは、sは、Vec<String> 型で文字が入っている状態になる。
use proconio::input;
fn main() {
input! {
n: i32,
s: [String; n]
}
println!("{:?}", n);
println!("{:?}", s);
}
$ cat hoge.txt | cargo run
3
["hello", "world", "atcoder"]
cagro-compete を使う
今まで、標準ライブラリを使う方法と、proconio を使う方法などを確認してきたが、問題ごとに cargo プロジェクトを作って、cargo.toml を編集したり、テストデータを用意するのはなかなか骨が折れる作業となる。
これを緩和し、迅速に問題を解くための環境を用意する方法として、cargo-compate という cargo サブコマンドを導入する。
インストール
自分の環境の場合、これを行わないとインストールが成功しなかったため、Rust を現在の最新バージョンである 1.15.0 に更新する。
rustup upgrade
その後、以下のコマンドでインストールが完了する。
cargo install cargo-compete
初期セットアップ
AtCoder で利用するので、以下のコマンドを実行する。
mkdir atcoder
cd atcoder
cargo compete init atcoder
そうすると、以下のように聞かれる。proconio などが揃っている状態を使いたいので、2を選ぶ。
Do you use crates on AtCoder?
1 No
2 Yes
3 Yes, but I submit base64-encoded program
これにより、ファイルがいくつか作られる。AtCoderにログインする。
IDとパスワードが聞かれるので、入力する。
cargo compete login atcoder
ABC120用のパッケージを作成する。これはコンテストごとに作成する。
今回は、ABC120 を例に取る。
cargo compete new abc120
abc120 というフォルダが作成されるので、そのフォルダをお好みのIDEなどで開く。
(atcoder フォルダを開いても良いのだが、IntelliJ が、Cargo.toml を最初から認識しなかったりしたので、abc120 を開くようにしました。)
問題を読みたい！と思ったら、以下のようにすることで、ブラウザでいっきに開くことができる。
便利。
cd abc120
cargo compete open
ということで、試しにA問題を解いてみる。
高橋くんは、自動販売機でジュースを買ったときの音が好きです。
その音は1回A円で聞くことができます。
高橋くんは B 円持っていますが、お気に入りの音を C 回聞くと満足するため、B円で最大 C回まで聞けるだけ聞きます。
高橋くんはお気に入りの音を何回聞くことになるでしょうか。
入力は
A B C
でそれぞれ、1から100までの整数とのこと。
出力は整数を1つ出力するだけ。
単純な問題で、B円 / A円 (小数点切捨て) の回数と、Cを比べて B / A > C であれば、C を、そうでなければ B / A を出せば良いということになる。
ということで書いたのが以下
use proconio::input;
fn main() {
input! {
a: u32,
b: u32,
c: u32
}
let d = b / a;
let ans = if d > c { c } else { d };
println!("{}", ans);
}
テストをする場合は、abc120のフォルダで 以下を実行する。
a 問題のテストをするという意味になる。
cargo compete test a
これにより、以下のようなテスト結果が出る。
1/3 ("sample1") Accepted (264 ms)
2/3 ("sample2") Accepted (264 ms)
3/3 ("sample3") Accepted (263 ms)
1/3 ("sample1") Accepted (264 ms)
stdin:
2 11 4
expected:
4
actual:
4
2/3 ("sample2") Accepted (264 ms)
stdin:
3 9 5
expected:
3
actual:
3
3/3 ("sample3") Accepted (263 ms)
stdin:
100 1 10
expected:
0
actual:
0
ちなみに、間違えると、以下のようになる。(条件演算子を逆にした場合)
1/3 ("sample1") Wrong Answer (465 ms)
2/3 ("sample2") Wrong Answer (465 ms)
3/3 ("sample3") Wrong Answer (465 ms)
1/3 ("sample1") Wrong Answer (465 ms)
stdin:
2 11 4
expected:
4
actual:
5
2/3 ("sample2") Wrong Answer (465 ms)
stdin:
3 9 5
expected:
3
actual:
5
3/3 ("sample3") Wrong Answer (465 ms)
stdin:
100 1 10
expected:
0
actual:
10
error: 3/3 tests failed
提出したい場合は、以下のコマンドを実行する。これも、abc120のフォルダで 実行する必要がある。
cargo compete submit a
1/3 ("sample1") Accepted (6 ms)
2/3 ("sample2") Accepted (3 ms)
3/3 ("sample3") Accepted (5 ms)
1/3 ("sample1") Accepted (6 ms)
stdin:
2 11 4
expected:
4
actual:
4
2/3 ("sample2") Accepted (3 ms)
stdin:
3 9 5
expected:
3
actual:
3
3/3 ("sample3") Accepted (5 ms)
stdin:
100 1 10
expected:
0
actual:
0
┌───────────────────┬─────────────────────────────────────────────────────────┐
│ Method │ cargo-compete │
├───────────────────┼─────────────────────────────────────────────────────────┤
│ Language ID │ 4050 │
├───────────────────┼─────────────────────────────────────────────────────────┤
│ Size │ 188 │
├───────────────────┼─────────────────────────────────────────────────────────┤
│ URL (submissions) │ https://atcoder.jp/contests/abc120/submissions/me │
├───────────────────┼─────────────────────────────────────────────────────────┤
│ URL (detail) │ https://atcoder.jp/contests/abc120/submissions/20514432 │
└───────────────────┴─────────────────────────────────────────────────────────┘
Successfully submitted the code
│ 2021-02-27 18:37:36 +09:00 │ A - Favorite Sound │ Rust (1.42.0) │ AC │ 4 ms │ 2132 KB │
結果が上記のように表示される。AC が出れば成功。嬉しい。
ooharabucyouABC120 - B問題を解く
ということで、引き続きB問題をときます。
正整数A, Bが与えられます。
AもBも割り切る正整数のうち、K番目に大きいものを求めてください。
なお、与えられる入力では、割り切る正整数のうち K番目に大きいものが存在することが保証されます。
入力
A B K
- 入力はすべて整数である
- 1 ≦ A, B ≧ 100
- AもBも割り切る正整数のうち、K番目に大きいものが存在する。
- K ≧ 1
--
ということで、解き方を考える。A, B もたかだか100までしか数がないので、100から1までA, B が割り切れるかを検証していき、K番目に割り切れる数を出力してプログラムを止めればよい。
が、より効率的に求めるために、以下を考える。
- 最大の割り切れる数は A, B のどちらか小さい数より大きな数はありえない
ということろで、100から数を見る必要はなく、MIN(A, B) から 1 までで検証すればよい。
さっと書いたのが以下。
use proconio::input;
fn main() {
input! {
a: u32,
b: u32,
k: u32
}
let mut c = 0;
for i in (1..=std::cmp::min(a, b)).rev() {
if a % i == 0 && b % i == 0 {
c += 1;
if k == c {
println!("{}", i);
break;
}
}
}
}
std::cmp::min は、標準ライブラリで用意されている、2つの値の最小値を求めるもの。
1..=std::cmp::min(a, b) とすることで、1から、A, B の最小値までの Range 型になる。これは for により、範囲を1つつづ取り出すことができる。
ここで 1..std::cmp::min(a, b) としてしまうと、1から、A, Bの最小値の-1までしか範囲にならないので注意したい。
今回は、「1から順番に」ではなく、「MIN(A,B)から1まで順番に」なので、Range 型など for により取り出せる型にトレイトにて実装されている rev() を利用して、逆にする。
あとは、a でも、b でも割れるかを判定し、何番目かどうかを管理する変数をカウントアップしていき、K番目になったら出力して終了！という感じ。
cargo compete submit b
1/3 ("sample1") Accepted (3 ms)
2/3 ("sample2") Accepted (5 ms)
3/3 ("sample3") Accepted (6 ms)
1/3 ("sample1") Accepted (3 ms)
stdin:
8 12 2
expected:
2
actual:
2
2/3 ("sample2") Accepted (5 ms)
stdin:
100 50 4
expected:
5
actual:
5
3/3 ("sample3") Accepted (6 ms)
stdin:
1 1 1
expected:
1
actual:
1
┌───────────────────┬─────────────────────────────────────────────────────────┐
│ Method │ cargo-compete │
├───────────────────┼─────────────────────────────────────────────────────────┤
│ Language ID │ 4050 │
├───────────────────┼─────────────────────────────────────────────────────────┤
│ Size │ 338 │
├───────────────────┼─────────────────────────────────────────────────────────┤
│ URL (submissions) │ https://atcoder.jp/contests/abc120/submissions/me │
├───────────────────┼─────────────────────────────────────────────────────────┤
│ URL (detail) │ https://atcoder.jp/contests/abc120/submissions/20529966 │
└───────────────────┴─────────────────────────────────────────────────────────┘
Successfully submitted the code
│ 2021-02-27 21:13:29 +09:00 │ B - K-th Common Divisor │ Rust (1.42.0) │ AC │ 6 ms │ 2124 KB │
ということで無事提出。
ooharabucyouABC120 - C問題を解く
ということで、C問題です。
0と1で構成される文字列を読み取って、特定の処理は何回することができるか? というもの。
下手に文意通り処理するのではなく、0と1の数を数えるだけでよさそう。
0と1が1個づつであれば、2個消せる = 01 or 10
0が2個、1が1個であれば、2個消せる = 001, 010, 100
0が2個、1が2個であれば、4個消せる = 0011, 0110, 1100, 1010, 0101, 1001... どれを選んでも4個かならず消せる
0が2個、1が10個の場合、4個消せる = 011111111110, 1001111111111... どれをとっても4個消せる
0の個数と、1の個数を取り、最小値を求めて2倍する、という感じで良さそうだ。
use proconio::input;
use std::cmp::min;
fn main() {
input! {
s: String
}
let a = s.match_indices("0").count();
let b = s.match_indices("1").count();
println!("{}", min(a, b) * 2);
}
match_indices(文字列) は、指定した文字列が存在する箇所を MatchIndices<&str> で返してくれる。
これは、繰り返し取得できる trait が実装されている。 (Iterator) ので、count() によりヒットした件数を返してくれる。
cargo compete submit c
1/3 ("sample1") Accepted (6 ms)
2/3 ("sample2") Accepted (3 ms)
3/3 ("sample3") Accepted (4 ms)
1/3 ("sample1") Accepted (6 ms)
stdin:
0011
expected:
4
actual:
4
2/3 ("sample2") Accepted (3 ms)
stdin:
11011010001011
expected:
12
actual:
12
3/3 ("sample3") Accepted (4 ms)
stdin:
0
expected:
0
actual:
0
┌───────────────────┬─────────────────────────────────────────────────────────┐
│ Method │ cargo-compete │
├───────────────────┼─────────────────────────────────────────────────────────┤
│ Language ID │ 4050 │
├───────────────────┼─────────────────────────────────────────────────────────┤
│ Size │ 211 │
├───────────────────┼─────────────────────────────────────────────────────────┤
│ URL (submissions) │ https://atcoder.jp/contests/abc120/submissions/me │
├───────────────────┼─────────────────────────────────────────────────────────┤
│ URL (detail) │ https://atcoder.jp/contests/abc120/submissions/20544406 │
└───────────────────┴─────────────────────────────────────────────────────────┘
Successfully submitted the code
│ 2021-02-27 22:04:15 +09:00 │ C - Unification │ Rust (1.42.0) │ AC │ 11 ms │ 2288 KB
proconio の憂鬱
IntelliJ だと、マクロで作られた変数は入力補完してくれないのであった。
これは、どうやら Cargo.toml に、proconio と、同じく入力を補助してくれる whiteread の両方が依存として設定されているとこの問題が起きるらしい。
proconio を使っている場合、whiteread を使う必要はなさそうなので、コメントアウトしてしまう。
- whiteread = "=0.5.0"
+ # whiteread = "=0.5.0"
これにより、無事に入力補完されるようになった。
今後、cargo compete new するときに、これを同じように対処してほしいので、init 時に作られたファイルである、compete.toml を少し編集する。
+whiteread = "=0.5.0"
-#whiteread = "=0.5.0"
ABC120 - D問題を解く
ということであっという間にD問題。
橋がつながっている島の橋を1本ずつ外していったときに、全くつながっていない島同士の組み合わせの数を求めていく。というものですが、「答えが32bit整数に納まらない」というところを見るに、そのままの形で問題を実装しないほうが良さそう。
最終的にはすべての橋が崩落して、どこの島にもつながっていない状態から始めていくのが良さそうである。4つの島がある場合は、
そこから橋を1本づつつなげていくのであるが、どっかとつながっているものと、全くつながっていないものを管理するときにグループという考え方を使えれば良さそうだ。
4つの島があるとき、最初の段階だとそれぞれ独立した状況である。
不便さは、
1
2
3
4
最初に、1-4でつなげたとき、以下のようなグループになる。
1つの独立した島と、1つの独立した島がつながって、1つの組み合わせが成立したので、初期の不便さから1が引かれて、5となる。
1 4
2
3
次に、2-3でつなげたときに、以下のようなグループになる。
先程と同じく、1つの独立した島どうしが繋がり、不便さは
1 4
2 3
さらに、1-3でつなげたときに、以下のようなグループになる。2つの島のグループどうしが繋がり、2 * 2 で、4の組み合わせが解消された。不便さは
1 2 3 4
この手のグループを管理するとき、自分がどのグループなのかを一々検索していたりなどすると、おそらく時間切れになってしまう。
高速に検索するためには、UnionFind を使うと良さげだ。
UnionFind 構造体を作り管理だ！と考えても良いが、ここは便利なことにAtCoderでは、petgraph という Cargo クレートが使えるようになっているらしい。
と思ったのだが、今回のケースだと、それぞれのグループにどれだけのノードがあるのかということを配慮に入れる必要があるので自前で実装する。
ということで書いたのが以下。
C問題に比べるとだいぶコード量増えている。。
- 配列について、0番地から扱うのが面倒そうなので、
n + 1のサイズの配列で取り扱い、0番地は使わない。 - それぞれの木で管理しているノードの数のための
tree_sizeを用意する。 - Union するときに、木のノード数に関する処理を行う。
- Vec を逆にした状態で for でループしたい場合は、一度 iter() により、
Iterにした上で、rev()をかける。 - 答えが32bit整数を超える可能性もあるとのことだが、64bitなマシンで動いていると思われるので、
usizeを使っておけば問題がない。
use proconio::input;
struct UnionFind {
root: Vec<usize>,
tree_size: Vec<usize>,
count: usize
}
impl UnionFind {
fn new(size: usize) -> UnionFind {
UnionFind {
root: vec![0; size + 1],
tree_size: vec![1; size + 1],
count: size * (size - 1) / 2
}
}
fn union(&mut self, a: usize, b: usize) {
let fa = self.find(a);
let fb = self.find(b);
self.unite(fa, fb);
}
fn find(&mut self, a: usize) -> usize {
if self.root[a] == 0 {
return a;
} else {
self.root[a] = self.find(self.root[a]);
return self.root[a]
}
}
fn unite(&mut self, a: usize, b: usize) -> () {
if a == b {
return ()
}
self.root[b] = a;
self.count -= self.tree_size[a] * self.tree_size[b];
self.tree_size[a] += self.tree_size[b];
self.tree_size[b] = 0;
}
}
fn main() {
input! {
n: usize,
m: usize,
links: [(usize, usize); m]
}
let mut fuben: Vec<usize> = Vec::new();
let mut uf = UnionFind::new(n);
// 橋がまったくない状態から橋を1本ずつかけて、
// 不便さを求める
for link in links.iter().rev() {
fuben.push(uf.count);
uf.union(link.0, link.1);
}
// 不便さを逆に出す
for f in fuben.iter().rev() {
println!("{}", f);
}
}
一通りやってみたので、次はWebアプリケーションの開発をやってみる。