こんにちは、Progateのキリル(@virtualkirill)です。
本記事は Progate Advent Calendar 16日目の記事になります。
ビンゴというゲームをご存知ですか？知っている人が多いと思いますが、5x5 のカード（ボード）に並べたマスに書かれた番号と、順番に発表される番号がマッチしたら、そのマスをマークして、一列（縦あるいは横）がマークできたら勝ちというシンプルなルールです。
今回は、発表用の番号と、複数のカードを与えるとどのカードが先に勝つかを計算するプログラムを作っていきます。これは毎年行われるAdvent Of Codeのパズルの一つで、このパズルの実装に特に向いている Rust を選びました。
インプット
インプットはこんな感じです。
7,4,9,5,11,17,23,2,0,14,21,24,10,16,13,6,15,25,12,22,18,20,8,19,3,26,1 22 13 17 11 0 8 2 23 4 24 21 9 14 16 7 6 10 3 18 5 1 12 20 15 19 3 15 0 2 22 9 18 13 17 5 19 8 7 25 23 20 11 10 24 4 14 21 16 12 6 14 21 17 24 4 10 16 15 9 19 18 8 23 26 20 22 11 13 6 5 2 0 12 3 7
1 行目は発表される番号で、その下にはスペースで区切られたボードがたくさんあります。
アプローチ
まずは、インプットを発表番号とそれ以外にわけて、各ボードをu8の配列としてパースします。u8にしたのは、どの番号も 100 を超えることなく 8 ビットに収まりそうだからです。
fn main() { let input = include_str!("input.txt"); let mut lines = input.split("\n\n"); let drawn: Vec<u8> = lines .next() .unwrap() .split(",") .map(|ch| ch.parse().unwrap()) .collect(); }
splitメソッドが返すのはIteratorです。タイトルでは Iterator をたくさん使うことを約束しましたが、ここが出番です。Iterator トレイトには next メソッドがあり、そのトレイトを実装するなんらかのコレクションのアイテムをステップ・バイ・ステップに取得することができます。Iterator 自体はインターフェースだったり、プロトコルだったりさまざま名前でさまざまな言語に存在していて馴染んでいる人もいるかと思いますが、Rust ほど Iterator を実装したスタンダードライブラリの要素を持つ言語はなかなかないです。
Rust の Iterator は lazy です。というと、実際に処理のチェーンの結果を求めるまで、なにも計算されることがない。新しいコレクションに結果を保存するためには collect する必要があります。collect してはじめて計算がはしり十分なメモリがアロケートされます。そのためたくさんの処理をチェーンして、途中の状態を気にしなくてよくなります。 collect はかなり抽象的で、たまにはコレクションの型を指定する必要があります。上記では、発表番号の字をu8として parse していますが、何型として parse するか指定しないと collect が混乱します。そのためVec<u8>というふうに指定しました（一方で、そこで指定すると parse に型を教えなくても推論が効きます）。
次はボードです。Primitive obsession (基本データ型への執着)というフレーズがありますが、なんでもかんでも基本データ型で表そうとすることをさしています。ビンゴのボードの番号はマークされているかされていないかの二つの状態があるので、ボードと同じサイズの、0 と 1 しか入れない配列を作って状態をトラッキングしてもいいですが、それだと「基本データ型への執着」が強いです。なので Board という概念の定義からはじめます:
use std::collections::HashMap; type State = HashMap<u8, bool>; #[derive(Debug, Clone)] struct Board { state: State, numbers: Vec<u8>, width: usize, } impl Board { fn new(state: State, numbers: Vec<u8>, width: usize) -> Self { Self { state, numbers, width, } } }
十分わかりやすいでしょう。要件を見てみると、特定の番号をボードで探す処理が多くなりそうなので、O(n)の find を避けるために、番号の値で状態がわかるように HashMap に初期化します。HashMapだと与えられた番号でO(1)でその状態を取得することができます。
let mut boards = lines .map(|board| { let numbers: Vec<u8> = board .split_ascii_whitespace() .map(|num| num.parse().unwrap()) .collect(); let state = numbers.iter().map(|num| (*num, false)).collect(); Board::new(state, numbers, 5) }) .collect();
collect がたくさんありますね。まずはボードの番号を普通の Vec として parse し collect します。次は初期状態としてすべての番号はマークされていないので、その番号の Vec を（番号、false）タプルにマップするのですが、なんと collect で HashMap が作れるので非常に便利です。そして最後に結果をボードの Vec に collect します。
パズルの要件では、一番最初に勝つボードを見つけるだけではなく、その勝ったボードのマークされていないすべての番号の和と勝ったときの番号の積を答えないといけないです。それを達成するには、以下のような処理が必要になりそうです。
if let Some((board, num)) = find_winning_board(&drawn, &mut boards) { let sum = board.sum_unmarked(); println!("{}", num as usize * sum) }
find_winning_board がメインです。その中身をどうすれば答えにたどり着くかイメージしてみましょう。
fn find_winning_board(drawn: &Vec<u8>, boards: &mut Vec<Board>) -> Option<(Board, u8)> { for n in drawn { for board in boards.iter_mut() { if let Ok(_) = board.try_complete(*n) { return Some((board.clone(), *n)); } } } return None; }
イメージといいながら答えを半分出しています。もちろん投稿する前に解いてしまっているのでそのままコピペしています。やることは非常にシンプルです。発表番号をひとつずつ取り、全部のボードでその番号を探し、見つかればマークします。ついでにそれでボードの一列（縦か横）の番号がマークされれば勝ちなので早期 return をします。iter_mutは、コレクションをIteratorにするメソッドiter()のミュータブル版です。それぞれのボードの番号の状態を変更しないといけないため、ボードへのミュータブルな参照を取得します。
try_completeでは、番号をマークした後に、ボードの全ての例の中で、全てマークされている列がないかを確認します。どういうふうに実装するのがいいかな？
僕がたどり着いた答えはこんな感じです。
fn try_complete(&mut self, num: u8) -> Result<(), ()> { if let Some(checked) = self.state.get_mut(&num) { *checked = true; for i in 0..self.width { let mut row = self.numbers.iter().skip(self.width * i).take(self.width); let mut col = self.numbers.iter().skip(i).step_by(self.width); if row.all(|n| self.state[n]) { return Ok(()); } if col.all(|n| self.state[n]) { return Ok(()); } } } return Err(()); }
まずは発表番号がそもそもボードに存在しているかどうかを見ます。もし存在していたらついでにそれをマークします。次はまたIteratorの出番です。マーク状態を確認するとき、番号単位ではなく、列単位でみたい。これはいくつかのやり方があると思いますが、一つのやり方としては以下のような順番で番号の列を交互に全て確認します（5x5の場合は全部で10回）。
x x x x x . . . . . . . . . . . . . . . . . . . . x . . . . x . . . . x . . . . x . . . . x . . . . . . . . . x x x x x . . . . . . . . . . . . . . . . x . . . . x . . . . x . . . . x . . . . x . . . . . . . . . . . . . x x x x x . . . . . . . . . .
このイメージをコードに直接訳す前に、人間の言葉でざっと命令にするとどうなるでしょうか？
- 最初の横の列（５つの番号）を確認する
- 縦の列は、ボードの幅と同じ個数の番号をステップ幅として全て確認する
- 次の横の列は、ボードの幅と同じ個数の番号をスキップして、そこから５つの番号を確認する
- 次の縦の列は、前回の起点の次の番号を起点として、またボードの幅と同じ個数の番号をステップ幅として全て確認する
- ...
アルゴリズムができそうです。これをそのままコードに訳すと、ネストした複雑なループになってしまうけど、もう一度重要な部分をコピペします：
for i in 0..self.width { let mut row = self.numbers.iter().skip(self.width * i).take(self.width); let mut col = self.numbers.iter().skip(i).step_by(self.width); }
iter()の結果には、さまざま便利なメソッドが実装されています。人間の言葉でアルゴリズムを書いた時は「スキップ」という単語を使ったが、やろうとしていることをまさに表すメソッドskipがあります。横の列の場合は、0..ボードの幅の範囲で、スキップする番号の数を
Iteratorに教えます（0, 5, 10, 15...)。skipだけだとスキップはするけど、最後まで取得してしまうので、takeで確認したい番号だけの数を指定できます。
同じく縦の列の場合は、0..ボードの幅の範囲で、起点を1ずつずらして、ステップ幅をstep_byという便利なメソッドで指定します。そして最後にallというメソッドと、ボードの番号→状態のマップを使って列の状態を確認します。
if row.all(|n| self.state[n]) { return Ok(()); } if col.all(|n| self.state[n]) { return Ok(()); }
これでほぼ終わりです。パズルの答えに必要な和、上記で出てきた board.sum_unmarked()はどう実装するか？もちろん Iteratorを使います。
fn sum_unmarked(&self) -> usize { self.numbers .iter() .filter(|n| !self.state[n]) .map(|n| *n as usize) .sum() }
番号のベクタであるVec<u8>をIteratorにすると、sumという自明で便利なメソッドが使えます。（sumは u8にはできなくて一回usizeに変換する必要がありました）。
パフォーマンス
パズルなのでそこまでパフォーマンスは気にしなくてもいいと思いますが、一点注目したい点があります。try_completeで使うIteratorはどれもcollectしていない。なぜかというと、最終的にallで状態をしたいだけなので、列を別のところで再利用するつもりがないからです。これだと、毎回一つの独立したコレクションとしてそれぞれの列を扱っているように見えても本当はnumbersの中身しかアクセスしていないです。新しいVecはヒープにアロケートされますが、ヒープのアロケーションがプログラムが遅くなる原因の一つです。上記のコードで collectなどでアロケートをしてみると実行が2倍ぐらい遅くなってしまいます！しかし上記でメンションしたようにRustのIteratorはlazyなので不要なアロケーションを回避できて素晴らしいと思います。
フルコード
use std::collections::HashMap; type State = HashMap<u8, bool>; #[derive(Debug, Clone)] struct Board { state: State, numbers: Vec<u8>, width: usize, } impl Board { fn new(state: State, numbers: Vec<u8>, width: usize) -> Self { Self { state, numbers, width, } } fn sum_unmarked(&self) -> usize { self.numbers .iter() .filter(|n| !self.state[n]) .map(|n| *n as usize) .sum() } fn try_complete(&mut self, num: u8) -> Result<(), ()> { if let Some(checked) = self.state.get_mut(&num) { *checked = true; for i in 0..self.width { let mut row = self.numbers.iter().skip(self.width * i).take(self.width); let mut col = self.numbers.iter().skip(i).step_by(self.width); if row.all(|n| self.state[n]) { return Ok(()); } if col.all(|n| self.state[n]) { return Ok(()); } } } return Err(()); } } fn find_winning_board(drawn: &Vec<u8>, boards: &mut Vec<Board>) -> Option<(Board, u8)> { for n in drawn { for board in boards.iter_mut() { if let Ok(_) = board.try_complete(*n) { return Some((board.clone(), *n)); } } } return None; } fn main() { let input = include_str!("input.txt"); let mut lines = input.split("\n\n"); let drawn: Vec<u8> = lines .next() .unwrap() .split(",") .map(|ch| ch.parse().unwrap()) .collect(); let mut boards: Vec<Board> = lines .map(|board| { let numbers: Vec<u8> = board .split_ascii_whitespace() .map(|num| num.parse().unwrap()) .collect(); let state = numbers.iter().map(|num| (*num, false)).collect(); Board::new(state, numbers, 5) }) .collect(); if let Some((board, num)) = find_winning_board(&drawn, &mut boards) { let sum = board.sum_unmarked(); println!("{}", num as usize * sum) } }
