rust
競技プログラミング

AtCoder に登録したら解くべき精選過去問 10 問を Rust で解いてみた

はじめに

drken さんの素晴らしい記事で紹介されていた AtCoder に登録したら解くべき精選 10 問 を、お気に入りの言語の Rust で解いてみました。
扱う問題一覧はこちら https://beta.atcoder.jp/contests/abs/tasks

自分が Rust で競技を始めたときはかなり躓いてドキュメントや stackoverflow を見まくったので、そんな負担を少しでも軽減できればと思い書きました。

対象

この記事は入門書をさらっと読んだ方、または他の言語経験があって Rust の雰囲気だけでも知りたい人が対象です。

入出力

どの言語でも最初に躓くのが入力だと思います。Rust も例にもれずそうなりがちで、Rust 競プロerがそれぞれ独自に入出力用ツールを自作しているのが現状です。
この記事では、一例として次の入力関数を使います。以降の記事中のソースコードは定義を省略します。
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 マクロを使えばあまり難しくありません。簡単な使い方は以降のコードを見たほうが早い気がします。
println も、高速ではありません。

目安として、現在 (2018/09/24) の AtCoder で $10^5$ 個オーダーの整数を read で入力すると数十ミリ秒、println で出力すると数百ミリ秒かかります。

2018/04/02 追記 :
tatsuya6502 さんが入出力高速化についてフォローしてくださいました。ありがとうございます!
Rustの競プロ向け入力関数やマクロを高速化できるかやってみた

2018/09/24 追記 :
tanakh 先生による便利マクロです。
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 文を評価することができます (したがって文より式と呼んだほうがいい場面もあります) 。

変数 aread() で読み取った値を代入する文は、この場合は以下のどれでもコンパイルが通ります。

let a: u32 = read(); // (1)
let a = read::<u32>(); // (2)
let a: u32 = read::<u32>(); // (3)

(1), (2) では、最も冗長な (3) で重複している u32 型の明示のうち片方を省略しています。(1) では、read の定義における T が、(2) では a の型が推論されるためこれで valid です。以降に登場しますが、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 の文字列の概念は少し複雑で、他の言語と異なるところが多いです。その一例としてインデックスで文字にアクセスができません。

なので、ここでは chars メソッドで文字を舐めるイテレータを生成し、filter1 だけを残し、残った個数を出力しています。
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> に変換し、直接インデックスでアクセスして判定しています。
このように、競技プログラミングで文字列の扱いに困ったら、Vec<char> または Vec<u8> に変換するのも手です。
Vec<char> に変換するには、 chars メソッドで Iterator<Item = char> を取得し、collect() します。String 型には into_bytes という Vec<u8> に変換するメソッドがあります。

collect は、Vec を始めとする std::collections::VecDequestd::collections::BTreeSet などのデータ構造を構築するメソッドです。s の型が Vec<char> であると明示すると、collect に型引数を渡さなくても Vec<char> 型と推論してくれます。

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..nmap で割り切れる回数に写し、最小値を求める方針にします。map に渡すクロージャで u32 型の整数を読み込み、さらに、整数型には trailing_zeros という便利なメソッドがあるので、それを使って 1 の桁から連続する 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 ループを回します (外部ライブラリにはありますが、それをダウンロードしてこれるオンラインジャッジは今のところありません)。
もっと良い書き方あるのかなあ。

第 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 は Unicode なども考慮した「文字」なので引き算ができません。

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 < xy == xy > x かを判定し、対応する列挙型 Ordering のメンバを返します (ドキュメント参照)。reverse() によって結果を反転し、sort_by によってそれをもとにソートします。y.cmp(x) でも同じです。

enumerate は Ruby の Enumerate#each_with_index や、Python の zip(hoge, range(n)) に対応するもので、インデックスと要素への参照のタプルのイテレータを作成します。

-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)$ の安定ソートアルゴリズムであることが保証されており、現バージョン 1.29.0 ではナイーブなマージソートが採用されています。
より高速にソートを行いたく、安定でなくてもよい場合は sort_unstable を使います。(ただし、残念ながら現在 (2018/03/14) の AtCoder のバージョンでは使えないようです。。。)

第 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 メソッドを使うのもきれいですし、メモリ的にはこちらの方が効率的です。

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);
}

解説

多重ループを使います。Rust には当然 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入門

まず、扱いやすいように、全ての文字列に対し、charsrev で反転してから Vec<char> にして s に代入します。
let mut s = &s[..] は、s の最初から最後までの要素を slice (C++ でいうところの長さ固定の配列) として取得し、その参照への mutable 変数を取得しています。
何を参照しているかが mutable なのであって、参照されている値そのものは immutable のままです。

ここで s が二度定義されています。このようにすると定義が上書きされます (シャドーイングといいます)。 Vec<char> 型の s は見えなくなるだけで、リソースはスコープを抜けるまで有効であり、動的型付け言語のように型を上書きしてしまっている訳でもありません。

while 文の中で、 slicestarts_with というメソッドを使って、パターン (patterns) の中に prefix が一致するものが無いか調べています。
find の戻り値は Option 型です。上記の条件を満たす要素が存在すれば、最も最初に見つかったものを Some で包んで返し、見つからなければ None を返します。
if let Some(p) = ... と書くと、もし戻り値が Some なら s = &s[p.len()..] とすることで s の先頭から s.len() 文字目から最後までの範囲の slice の参照で s を置き換えます。
存在しなかったら (else)、succeededfalse にして終了します。
s の長さが 0 になったら、空文字列は自明に置き換え可能なので、Yes/No を出力して終了します。

2018/03/14 22:17 追記 : 当初は計算量的によくない可能性があるコードを載せていましたが修正しました。
未検証ですが、見た目に反して最適化のおかげで計算量的に遅くならずに済むのかなぜか AC を得ていました。

2018/03/15 1:40 修正 : slicestarts_with というメソッドを使うコードに修正しました。

解答例 2

2018/03/15 12:28 追記 : Delta114514 さんの記事を見てもっとシンプルな方法を知りました。
よくよく考えると前から eraser -> erase -> dreamer -> dream の順に、String のメソッド replace を用いて空文字列に置換すればよいですね。こちらは思いつきませんでした。

そちらの解法で書き換えてみましょう。
さて、Rust のイテレータには次のループを initfunc に関して一般化するメソッド 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);

foldreplace を使ったシンプルな解答例が次のコードになります。

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 点)

問題概要

$n$ 個の時間と 2 次元平面上の座標の組 $(t_1, (x_1, y_1)), \ldots, (t_n, (x_n, y_n))$ が与えられます。便宜上、単位は秒とメートルとします。旅人は $1$ 秒に $1$ m だけ軸に平行な 4 方向のうちいずれかの向きに真っ直ぐ移動しなくてはいけません。$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 さんの記事で紹介されているように、ある点とその次に訪れる点の距離と時間の関係のみに注目すればいいです。slicewindows メソッドを使うと、隣り合う k 個の要素からなるブロックを前から順に走査する新しいイテレータを取得できます。all メソッドを使ってそれら全てが条件を満たすか判定できます。

おわりに

Rust はいいぞ