はじめに
drken さんの素晴らしい記事で紹介されていた AtCoder に登録したら解くべき精選 10 問 を、Rust で解いてみました。
扱う問題一覧はこちらです。
自分が Rust で競技プログラミングを始めたときはかなり躓いてドキュメントや stackoverflow を何度も調べたので、そのような負担を少しでも軽減できればと思い書きました。
対象
この記事は以下の入門書を軽く読んだ方、または他の言語経験があり Rust の雰囲気だけでも知りたい人が対象です。
The Rust Programming Language
入出力
どの言語でも最初に躓くのが入力だと思います。Rust も例にもれずそうなりがちですが、tanakh 先生による便利なマクロがおすすめです: Rustで競技プログラミングの入力をスッキリ記述するマクロ
この記事は上記の記事が書かれる前に書かれたので、以下のような入力関数 read を使います。以降の記事中のソースコードでは定義を省略します。
read は標準入力から空白をスキップしてトークンを受け取り、T 型に変換して返します。エラーハンドリングを行わない競技プログラミング用の簡素なものであり、読み込みと変換に失敗した場合は異常終了します。
use std::io::*;
use std::str::FromStr;
fn read<T: FromStr>() -> T {
let stdin = stdin();
let stdin = stdin.lock();
let token: String = stdin
.bytes()
.map(|c| c.expect("failed to read char") as char)
.skip_while(|c| c.is_whitespace())
.take_while(|c| !c.is_whitespace())
.collect();
token.parse().ok().expect("failed to parse token")
}
出力は、println マクロを使えばあまり難しくありません。簡単な使い方は以降のコードを見たほうが早い気がします。
read と println は少々遅いです。と言っても、それらがボトルネックになるような問題はあまり無いでしょう。
しかしたまにあるのも事実です。目安として、現在 (2018/09/24) の AtCoder で $10^5$ 個オーダーの整数を read で入力すると数十ミリ秒、println で出力すると数百ミリ秒かかります。このオーダーの時間が無視できないならより高速な入出力が必要です。そのような時は次の記事を参考にしてください: Rustの競プロ向け入力関数やマクロを高速化できるかやってみた
問題と解法の詳細について
問題文はごく簡単な文に要約して載せたので、厳密に知りたいときは AtCoder で全文を見てください。
また、解法もかなり簡単にしか説明しないので、詳細は drken さんの記事を参考にしてください。
第 1 問: ABC 086 A - Product (100 点)
問題概要
整数 $a$, $b$ の積が偶数か判定してください。
解答例
fn main() {
let a: u32 = read(); // (*)
let b: u32 = read();
let ans = if (a * b) % 2 == 0 { "Even" } else { "Odd" };
println!("{}", ans);
}
解説
3 項演算子はありません。しかし if ... else ... を評価することができます。そのような使い方をされているときは「if 文」よりも「if 式」と呼ぶ方が正確です。
コード中の (*) で示した行は u32 型の変数 a を read() で読み取った値で定義する文です。この場合は以下の全てに置き換えてもコンパイルが通ります。
let a: u32 = read(); // (1)
let a = read::<u32>(); // (2)
let a: u32 = read::<u32>(); // (3)
(3) が最も冗長で、(1), (2) では (3) で重複している u32 型の明示のうち片方を省略しています。(1) では、read の定義における T が : u32 から、(2) では <u32> から a の型が推論されるためこれで大丈夫です。後で紹介しますが、a の使われ方によって型が定まる場合は、単に let a = read() と書くことができます。
第 2 問: ABC 081 A - Placing Marbles (100 点)
問題概要
文字列 $s$ に 1 がいくつ含まれるか数えてください。
解答例
fn main() {
let s: String = read();
println!("{}", s.chars().filter(|&c| c == '1').count());
}
解説
Rust の文字列の概念は少し複雑で、文字列の種類によって複数の型が用意されています。他の言語ではよしなにしてくれる違いが厳密に分けられています。
その内の一つ String は添字で文字にアクセスができません。String は UTF-8 の文字列です。添字アクセスはマルチバイト文字を考慮すると前から順に何文字目か数える必要があるので線形時間の処理となり、考慮しないとマルチバイト文字の一部分だけ取得できるという気持ち悪い仕様になります。どちらも仕様としてイマイチなので最初から無しにするという思想なのでしょう。競技プログラミングなら ASCII コードの範囲に収まるよう入力フォーマットが配慮されているので余計なお世話ではありますが、仕方がありません。
上記の解答では、添字によるアクセスを行わない方法をとっています。String の chars メソッドで文字を舐めるイテレータ (型は Iterator<Item = char>) を生成し、filter で 1 だけを残し、残った個数を出力しています。char はマルチバイトも考慮した「文字」の型です。C++ などの四則演算ができる 1 バイトの「数値」である char 型とは根本的に異なると思ったほうがよいです。ちなみに、Rust におけるそれは u8 型です。
Iterator のメソッド filter は、bool 値を返すクロージャ (という大雑把に説明をするとその場で定義できる関数) を渡すと true となった要素だけを走査するイテレータに変換します。
文字列の扱い方を説明するために添字アクセスを行う解答も紹介します。
fn main() {
let s: Vec<char> = read::<String>().chars().collect(); // (*)
let mut cnt = 0;
if s[0] == '1' { cnt += 1; }
if s[1] == '1' { cnt += 1; }
if s[2] == '1' { cnt += 1; }
println!("{}", cnt);
}
ここでは、String を Vec<char> に変換し、添字でアクセスできるようにしてから実際に取得して判定しています。このように、競技プログラミングで String の扱いに困ったら、Vec<char> または Vec<u8> に変換するのも手です。Vec<char> はマルチバイトも考慮した「文字」の列、Vec<u8> は考慮しない「バイト (u8)」の列です。Vec<char> に変換するには、上のように chars メソッドで Iterator<Item = char> を取得し、collect メソッドを呼びます。
(*) で示した行の collect は、Vec を始めとする std::collections::VecDeque や std::collections::BTreeSet などのデータ構造を構築するメソッドです。s の型が Vec<char> であると明示すると、collect に型引数を渡さなくても Vec<char> 型と推論してくれます。以下の (1) はあえて明示せずに書き直したものです。また、(2) が示すように collect 側を明示すると let s 側は明示しなくてよくなります。どちらに置き換えてもコンパイルが通ります。
let s: Vec<char> = read::<String>().chars().collect::<Vec<char>>(); // (1)
let s = read::<String>().chars().collect::<Vec<char>>(); // (2)
Vec<u8> に変換するには String 型の into_bytes メソッドを使います。u8 値のリテラルは b'1' のように書きます。それによる解答例は以下です。
fn main() {
let s: Vec<u8> = read::<String>().into_bytes();
let mut cnt = 0;
if s[0] == b'1' { cnt += 1; }
if s[1] == b'1' { cnt += 1; }
if s[2] == b'1' { cnt += 1; }
println!("{}", cnt);
}
cnt の定義に mut が付いています。Rust では、mut が無い値は定数として、mut を付けると変数として定義されます。ここでは、後ろの行で cnt += 1 したいので、mut が必要です。
第 3 問: ABC 081 B - Shift Only (200 点)
問題概要
整数 $a_1, \ldots, a_n$ の中で $2$ で割り切れる回数が最小のものの、割り切れる回数を求めるてください。
解答例
fn main() {
let n = read();
println!(
"{}",
(0..n)
.map(|_| read::<u32>().trailing_zeros())
.min()
.unwrap()
);
}
解説
この問題は、よく考えると整数を 2 進数表記したときに、1 の位から連続する 0 の個数の最小値を求めればよいことが分かります。
今までは read で読み込む値の型を明示してきましたが、今回は明示していません。明示しなかった場合は、以降のその値の使われ方で推論できる場合は推論され、できなければエラーになります。下のコードでは (0..n) の部分から整数型であることが分かるので、明示する必要がありません。
Range オブジェクト 0..n を map で各々の割り切れる回数に写し、最小値を求める方針で解くことにします。map に渡すクロージャで u32 型の整数を読み込みます。さらに、整数型の trailing_zeros という最下位の桁から連続する 0 の個数を求める便利なメソッドで写し、最後に min を使ってその最小値を求めます。
min はイテレータの長さが 0 だった場合も考慮して、Option<T> 型を返す設計ですが、今回は $n > 0$ が保証されているので unwrap で剥がします。
第 4 問: ABC 087 B - Coins (200 点)
問題概要
整数 $a,b,c,n$ が与えられるので、$500i+100j+50k=x$ を満たす整数 $i,j,k \ (0 \le i \le a, 0 \le j \le b, 0 \le k \le c)$ の個数を求めてください。
解答例
fn main() {
let a: u32 = read();
let b: u32 = read();
let c: u32 = read();
let x: u32 = read();
let mut ans = 0;
for i in 0..a + 1 {
for j in 0..b + 1 {
for k in 0..c + 1 {
if i * 500 + j * 100 + k * 50 == x {
ans += 1;
}
}
}
}
println!("{}", ans);
}
解説
多重ループが必要です。残念ながら Ruby の Array#product のようなメソッドは標準ライブラリには無いので C++ と同様に for ループを回します (外部ライブラリにはありますが、それをダウンロードしてこれるオンラインジャッジは今のところありません)。もっと良い書き方あるのかなあ。
話が逸れるので詳しく解説しませんが、cargo-atcoder という任意のパッケージを利用してコードが書けるツールもあります。musl でシングルバイナリを生成し、base64 エンコードして提出するという仕組みです。
第 5 問: ABC 083 B - Some Sums (200 点)
問題概要
整数 $n,a,b$ が与えられます。$1 \le x \le n, a \le s(x) \le b$ を満たす $x$ に対し、$10$ 進数表記したときの桁の和の総和を求めてください。
解答例
fn main() {
let n: u32 = read();
let a: u32 = read();
let b: u32 = read();
let ans = (1..n + 1)
.filter(|x| {
let sum = x.to_string()
.chars()
.map(|c| (c as u8 - b'0') as u32)
.sum::<u32>();
a <= sum && sum <= b
})
.sum::<u32>();
println!("{}", ans);
}
解説
先程も登場した filter メソッドを使います。filter の中で整数から桁の和に変換し、条件で絞った後、sum メソッドで合計を求めます。
イテレータが内包するオブジェクトの型が分かっていても、sum に型引数を与えなければいけません。
整数の桁の和を求める方法はいろいろとありますが、上のコードでは to_string() で文字列に変換し、さらに各桁の文字 (char) を 0 から 9 の整数 (u8) に変換して和を求めます。上で説明したように char は UTF-8 の範囲まで考慮した「文字」なのでそのままでは引き算ができません。
a <= sum && sum <= b (セミコロンが無いことに注意) の部分で使っているように、関数とクロージャーは最後の文を式にした場合にそれが戻り値となります。
第 6 問: ABC 088 B - Card Game for Two (200 点)
問題概要
$a_1, \ldots, a_n$ を降順ソートし、偶数番目の和と奇数番目の和の差を求める問題です。
解答例
fn main() {
let n = read();
let mut a: Vec<u32> = (0..n).map(|_| read()).collect();
a.sort_by(|x, y| x.cmp(y).reverse());
let mut alice = 0;
let mut bob = 0;
for (i, &x) in a.iter().enumerate() {
if i % 2 == 0 {
alice += x;
} else {
bob += x;
}
}
println!("{}", alice - bob);
}
解説
Vec を昇順ソートするには sort メソッドを使えばよいのですが、今回は降順ソートしたいので、引数に関数またはクロージャをとる sort_by を使います。x > y でソートしたいときには |x, y| x.cmp(y).reverse() を与えればよいです。
cmp メソッドは、比較可能なオブジェクトを渡すと、比較して、y < x か y == x か y > x かを判定し、対応する列挙型 Ordering のメンバを返します (ドキュメント参照)。reverse() によって結果を反転し、sort_by によってそれをもとにソートします。y.cmp(x) でも同じです。
enumerate は Ruby の Enumerate#each_with_index や、Python の zip(range(n), *iterables) に相当するもので、インデックスと要素への参照のタプルのイテレータを作成します。
sort_by_key を使い -x で昇順ソートしてもいいです。ただし、unsigned 型にはマイナス記号が付けられなくなっているので、singed 型を使います。
let mut a: Vec<i32> = (0..n).map(|_| read()).collect();
a.sort_by_key(|&x| -x);
昇順ソート後に reverse してもいいです。
let mut a: Vec<u32> = (0..n).map(|_| read()).collect();
a.sort();
a.reverse();
ここで使った sort は最悪時間計算量 $O(n \log n)$ の安定ソートアルゴリズムであることが保証されており、ナイーブなマージソートが採用されています。
より高速にソートを行いたく、安定でなくてもよい場合は sort_unstable を使います。
第 7 問: ABC 085 B - Kagami Mochi (200 点)
問題概要
$a_1, \ldots, a_n$ から重複する要素を取り除いたときに何個残るか求めてください。
解答例
fn main() {
use std::collections::BTreeSet;
let n = read();
let a: BTreeSet<u32> = (0..n).map(|_| read()).collect();
println!("{}", a.len());
}
解説
この問題もいろいろな方法で解けますが、一番シンプルであろう方法が、collections::BTreeSet という集合を管理するデータ構造に突っ込んで重複を取り除くものでしょうか。
BTreeSet の他にも HashSet もあります。前者は内包する型が Ord トレイトの実装を持つこと (比較可能であること)、後者は Hash トレイトの実装を持つこと (ハッシュ値の計算ができること) が使える要件にあります。
dedup メソッドを使うのもきれいです。mut となってしまいますがメモリ的にはこちらの方が効率的です。
let mut a: Vec<u32> = (0..n).map(|_| read()).collect();
a.sort();
a.dedup();
println!("{}", a.len());
第 8 問: ABC 085 C - Otoshidama (300 点)
問題概要
整数 $n,x$ が与えられるので、$i+j+k=n, 10000i+5000j+1000k=x$ となる整数 $i,j,k$ を求めてください。
解答例
fn main() {
let n: i32 = read();
let x: i32 = read();
let mut ans = None;
'outer: for i in 0..n + 1 {
for j in 0..n - i + 1 {
let k = n - i - j;
if i * 10000 + j * 5000 + k * 1000 == x {
ans = Some((i, j, k));
break 'outer;
}
}
}
let (x, y, z) = ans.unwrap_or((-1, -1, -1));
println!("{} {} {}", x, y, z);
}
解説
多重ループを使います。当然 GOTO 文は存在しません。代わりに、ループのラベル付き break が使えます。下のコードでは一番外側のループに outer というラベルを付けて、答えが見つかったら抜けています。
C++ などでは解が見つかったことを表すフラグを用意するのがよいのかもしれませんが、Rust にはせっかく Option 型があるので使っています。
第 9 問: ABC 049 C - Daydream (300 点)
問題概要
文字列 $s$ が与えられます。dream, dreamer, erase, eraser を好きな順番で好きな個数だけ連結することで $s$ を作れるか判定してください。
解答例
fn main() {
let patterns: Vec<Vec<char>> = ["dream", "dreamer", "erase", "eraser"]
.iter()
.map(|s| s.chars().rev().collect())
.collect();
let s: Vec<char> = read::<String>().chars().rev().collect();
let mut s = &s[..];
let mut succeeded = true;
while s.len() > 0 {
let matched = patterns.iter().find(|&p| s.starts_with(p));
if let Some(p) = matched {
s = &s[p.len()..];
} else {
succeeded = false;
break;
}
}
println!("{}", if succeeded { "YES" } else { "NO" });
}
解説
解法は drken さんの記事で紹介されている貪欲法を使います。
上の解答例は、Rust の文字列に関するややこしい仕様をあえて避けないようなコードになっています。
文字列について全て説明すると長くなるのでここでは詳細を省略します。詳しくは次の記事を読んでください: Stringとstr入門
まず、扱いやすいように、全ての文字列に対し、chars を rev で反転してから Vec<char> にして s に代入します。
let mut s = &s[..] は、s の最初から最後までの区間を slice (C++ でいうところの長さ固定の配列) として取得し、その参照への mutable 変数を取得しています。
何を参照しているかが mutable なのであって、参照されている値そのものは immutable のままです。また、あくまで参照なので配列の一部のコピーは発生していません。C++ の std::string_view のようなものです。
ここで s が二度定義されています。このようにすると定義が上書きされます (シャドーイングといいます)。2 つめの s の定義以降で 1 つめの Vec<char> 型の s は見えなくなるだけで、リソースはスコープを抜けるまで有効であり、動的型付け言語のように別の型の変数で上書きしている訳でもありません。
while 文の中で、 slice 型が持つ starts_with というメソッドを使って、パターン (patterns) の中に prefix が一致するものが無いか調べています。
find の戻り値は Option 型です。上記の条件を満たす要素が存在すれば、最も最初に見つかったものを Some で包んで返し、見つからなければ None を返します。
if let Some(p) = matched と書くと、matched が Some なら真となり、s = &s[p.len()..] が実行されます。これは s の先頭から s.len() 文字目から最後までの範囲の slice の参照で s を置き換えます。
存在しなかったら (else)、succeeded を false にして終了します。
s の長さが 0 になったら、空文字列は自明に置き換え可能なので、Yes/No を出力して終了します。
解答例 2 (嘘)
2018/03/15 12:28 追記: Delta114514 さんの記事を参考に、よりシンプルな解放を追記しました。
2020/01/26 04:49 追記: Lily0727K さんの記事 によるその解法は間違いです！dreaerasem という反例が存在します。しかし fold の説明に使いたいので残しておきます。
よくよく考えると前から eraser -> erase -> dreamer -> dream の順に、String のメソッド replace を用いて空文字列に置換すればよいですね。こちらは思いつきませんでした。
そちらの解法で書き換えてみましょう。
さて、Rust のイテレータには次のループを初期値 init と関数 func に関して一般化するメソッド fold があります。
let mut acc = init;
for &x in iter {
acc = func(acc, x);
}
すなわち、上のコードと下のコードの実行後の acc の値は同じです。
let acc = iter.fold(init, func)
fold を使うと多くのループ処理が非常にシンプルに書けます。
例えば、0 から 9 までの整数の合計は次のようにすると求められます。
let sum = (0..10).fold(0, |acc, x| acc + x);
fold と replace を使ったシンプルな解答例が次のコードになります。
fn main() {
let patterns = ["eraser", "erase", "dreamer", "dream"];
let s = patterns.iter().fold(read::<String>(), |s, x| s.replace(x, ""));
println!("{}", if s.is_empty() { "YES" } else { "NO "});
}
第 10 問: ABC 086 C - Traveling (300 点)
問題概要
時間と 2 次元平面上の座標の $n$ 個の組 $(t_1, (x_1, y_1)), \ldots, (t_n, (x_n, y_n))$ が与えられます。便宜上、単位は秒とメートルとします。旅人は $1$ 秒に $1$ m だけ $x$ または $y$ 軸に平行な 4 方向のうちいずれかの向きに真っ直ぐ移動しなくてはいけません。つまり、1 秒に 1 度隣の格子点に移動します。$t=0$ のとき旅人は $(0, 0)$ にいます。$1$ 秒ごとに先の制約で自由に移動してよいとき、各 $i$ に対し時刻 $t_i$ に $(x_i, y_i)$ の位置にいることはできるか判定してください。
解答例
fn main() {
let n = read();
let mut v: Vec<(i32, i32, i32)> = (0..n).map(|_| (read(), read(), read())).collect();
v.insert(0, (0, 0, 0));
let yes = v[..].windows(2).all(|w| {
let (t, x, y) = w[0];
let (nt, nx, ny) = w[1];
let time = nt - t;
let dist = (nx - x).abs() + (ny - y).abs();
dist <= time && time % 2 == dist % 2
});
println!("{}", if yes { "Yes" } else { "No" });
}
解説
drken さんの記事で紹介されているように、ある点とその次に訪れる点の距離と時間の関係のみに注目すればいいです。slice の windows メソッドを使うと、隣り合う k 個の要素からなるブロックを前から順に走査する新しいイテレータを取得できます。all メソッドを使ってそれら全てが条件を満たすか判定できます。
おわりに
Rust はいいぞ。