Rustで競技プログラミングの入出力を自動的にやるマクロ

Rustで競技プログラミングの入力をスッキリ記述するマクロ の続編的なやつです。

モチベーション

AtCoderのRust環境も新しくなって、バージョンが新しくなったり一部の外部のcrateが使えるようになったり、ずいぶんいろいろできるようになりました。前回の記事では、コピペして使えるようなコンパクトで効果的なものを目指していましたが、そういう制約がない状況でならもっといろいろできるのではないかと考えていました。

もうちょっとだけ便利にできそうな気がしたので、とりあえずアイデアを形にしてみるかという感じで書いてみました。

アイデア

一般的に、人は標準入出力とやり取りをするのは苦手だけれども、関数の引数と返り値を扱うのは呼吸をするように容易に行えるものです。そして一般的に、競技プログラミングのタスクというのは入力が与えられて答えを出力するという自然に関数として表現できるものです。実際にTopCoderなどでは参加者が入出力のコードを書くのではなく、関数の引数と返り値で入出力をやり取りしています。

ところで以前作った input!マクロ は、宣言的に使えるものを目指した結果、構文が変数宣言とほぼ同じになりました。例えば、

fn solve() {
    input! {
        n: i64,
        m: i64,
    }
    ...
}

こういうのがあったとしたら、このinputの部分をそのまま引数に持っていけば、

fn solve(n: i64, m: i64) {
   ...
}

と、ただの関数になります。ということは、ただの関数から引数の部分を抜き取ってinput!マクロに入れてやると、引数で入力を受け取る関数から標準入力から入力を受け取る関数に簡単に変換できそうです。

返り値のほうはもっと簡単で、単にprintln!で出力すれば標準出力に出力する関数が出来上がります。

つまり、

fn solve(n: i64, m: i64) -> i64 {
    n + m
}

こういう関数を書いたら、自動的に標準入出力とやり取りするコードが生成できるんじゃないか？、というのがベースとなるアイデアです。

実装

実装してみたものが https://crates.io/crates/argio こちらになります(名前は Arguments to I/O converter 的な雰囲気で)。

#[argio]
fn main(n: i64, m: i64) -> i64 {
    n + m
}

と、普通の関数の頭に #[argio] という属性をつけると、その関数は引数と返り値ではなく、標準入出力からデータをやり取りする関数に変換されます。具体的には上記の関数が

fn main() {
    input! {
        n: i64,
        m: i64,
    }

    let ret = (|| -> i64 {
        n + m
    })();

    println!("{}", ret);
}

のように書き換わります。非常に単純なマクロです。

デフォルトでは引数を単に proconio::input! に渡しているので、関数の引数としては構文としておかしいものも受け取れます。例えば、整数のリストを受け取る次のようなものもうまく動きます（関数の引数に許されるものがどういうものかを厳密に把握しているわけではないのですが）。

#[argio]
fn main(n: usize, a: [i64; n]) -> i64 {
    a.into_iter().sum()
}

出力のカスタマイズ

出力は単に println! で行っているので、返り値が Display トレイトを実装している必要があります。これは、Vec<> などをそのままでは出力できないことを意味します。

配列を出力したいケースはよくあるので、どうやって解決しようかなと悩ましいところです。出力をDisplay トレイトではなく、このマクロが出力する用のトレイトを定義して、それを実装させる方針も考えられますが、今回は Display トレイトを実装するラッパを作って出力をカスタマイズする方式にすることにしました。例えば、次のように、

struct Wrap<T>(T);

impl<T: Display> Display for Wrap<Vec<T>> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for (ix, r) in self.0.iter().enumerate() {
            if ix > 0 {
                write!(f, " ")?;
            }
            r.fmt(f)?;
        }
        Ok(())
    }
}

#[argio]
fn main(n: usize) -> Wrap<Vec<usize>> {
    Wrap((0..n).map(|i| i * 2).collect())
}

Wrapといういわゆる "newtype" 型を定義して、これに対して Display を実装します。ここでは、Vec<T> をスペース区切りで表示するように定義したので、次のように動作します。

$ echo 10 | cargo run
0 2 4 6 8 10 12 14 16 18

しかしながらこんなのを毎回のんびり書いていたら競技プログラミングでは時間でぼろ負けするので、AtCoderでよく出てくる型に対するラッパーをあらかじめまとめて定義しました。

これを用いると、それをテンプレートに書くなりuseするなりして利用できます。

use competitive::AtCoder;

#[argio]
fn main(n: usize) -> AtCoder<Vec<usize>> {
    AtCoder((0..n).map(|i| i * 2).collect())
}

ただ、これでも毎回こんなのを書くのはめんどくさいので、問題ごとに書かなければならないコードの部分（関数の型と本体）ではなく、テンプレートの部分にラッパーの指定を押し込められるようにマクロのパラメーターで、使いたいラッパーを指定できるようにしました。

#[argio(output = AtCoder)]
fn main(n: usize) -> Vec<usize> {
    (0..n).map(|i| i * 2).collect()
}

output に指定したものでラップしてから println!するようになります。これで、関数の本体の部分からラッパーの指定が消えるので、問題固有のコード以外のノイズが消えたことになります。

配列をスペース区切りではなく改行区切りで出力したいような場合は、これもnewtypeラッパーを作ることになります。

struct Vertical<T>(Vec<T>);

impl<T> Display for AtCoder<Vertical<T>>
where
    T: Copy,
    AtCoder<T>: Display,
{
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        for i in 0..self.0 .0.len() {
            if i > 0 {
                writeln!(f)?;
            }
            AtCoder(self.0 .0[i]).fmt(f)?;
        }
        Ok(())
    }
}

#[argio(output = AtCoder)]
fn main(n: usize) -> Vertical<usize> {
    (0..n).map(|i| i * 2).collect_vec().into()
}

$ echo 5 | cargo run
0
2
4
6
8

出力を Display トレイトで行うのが難しい場合、もちろん手動で標準出力へ出力することもできます。その場合、関数の返り値はいらないし、何も表示してもらいたくはないはずです。関数の返り値を()にすれば、何も出力しない挙動になります。

#[argio(output = AtCoder)]
fn main(n: usize) {
    let ans = (0..n).map(|i| i * 2).collect_vec();
    for a in ans {
        println!("{}", a);
    }
}

$ echo 5 | cargo run
0
2
4
6
8

入力のカスタマイズ

出力のカスタマイズと同様に、入力もカスタマイズが考えられます。これは単にマクロを挿げ替えられるようにしました。

#[argio(input = proconio::input)]
fn main(n: usize) -> usisze {
    n * 2
}

inputに、関数の引数のトークン列を受け取って入出力を行うマクロを指定してやります。指定しなければproconio::inputが使われることになります。

マルチテストケース

argio は非常に単純なマクロですが、せっかくマクロにしたので、もうちょっと何かできないか考えてみます。プロコンでは先頭にケース数が来て、複数テストケースを処理させるような問題がちょくちょく出ます。そういうのを自動的に処理するコードは簡単にできそうなので対応してみました。

#[argio(multicase, output = Wrap)]
fn main(n: usize) -> Vec<usize> {
    (0..n).map(|i| i * 2).collect()
}

multicase というパラメーターを指定すると、

$ echo "3 2 3 5" | cargo run
Case #1: 0 2
Case #2: 0 2 4
Case #3: 0 2 4 6 8

こういう風にマルチテストケースを処理するコードが自動的に生成されます。各ケースの先頭にはケース番号を変数iとして、デフォルトで "Case #{i+1}: " が使われますが、指定することもできます。例えば、"" をしていすると、ヘッダが空文字になるので、

#[argio(multicase = "", output = Wrap)]
fn main(n: usize) -> Vec<usize> {
    (0..n).map(|i| i * 2).collect()
}

このようになります。

$ echo "3 2 3 5" | cargo run
0 2
0 2 4
0 2 4 6 8

フォーマットは https://crates.io/crates/interpol と大体同じような感じで指定できます。

副次的なメリット

Visual Studio Code の rust-analyzer 拡張機能は現在のところ、マクロ中で変数を定義した場合、その変数の型がうまく推論されないので、input!マクロで入力を行う場合、メソッドの補完や自動でアノテートされる変数の型がおかしかったりしていました。

その点このマクロであれば、見かけ上は単なる関数定義なので、変数の型を正しく扱ってくれます。

[(i64, i64); n] みたいな一見おかしな型でも、なんだかただのスライスとして扱ってくれているみたいです。

使用例

いくつか利用例を書いてみます。私が個人的に作っている competitiveというcrateを使っています。ここから argio をre-exportしています。

AtCoder Beginner Contest 174 A問題

あなたは、室温が30度以上のとき、またそのときに限り、冷房の電源を入れます。
今の室温は X 度です。冷房の電源を入れますか？
冷房の電源を入れるならば Yes、入れないならば No を出力せよ。

整数を受け取ってboolを返す関数です。AtCoderでは、boolを"Yes"と"No"で出力させることが大変多いので、AtCoder<bool> はそういう実装にしています。そうじゃなかったらboolではなく&strを返すことになると思います。

use competitive::prelude::*;

#[argio(output = AtCoder)]
fn main(x: i64) -> bool {
    x >= 30
}

構文上のシンタクティックノイズ（問題の本質と関係ない部分）をかなり消せたのではないかと思います。

AtCoder Beginner Contest 174 E問題

丸太が $N$ 本あり、それぞれ長さは $A_1, A_2, ..., A_N$ です。
これらの丸太を合計 $K$ 回まで切ることができます。長さ $L$ の丸太を片端から $t(0<t<L)$ の位置で切ると、長さ $t,L−t$ の丸太に分かれます。
丸太を合計 $K$ 回まで切った後最も長い丸太の長さが最小でいくつになるか求め、小数点以下を切り上げた値を出力してください。

n と k と a を入力するので、それぞれ main の引数に書いていきます。答えは非負整数になるのでusizeにします。

解き方は答えを直接二分探索すればよいです。切った後の丸太の長さの最大値を決めれば、各丸太に対して何回切らなければならないのかがすぐにわかるので、それがkより小さい範囲を探索すればよいという寸法です。

use competitive::prelude::*;

#[argio(output = AtCoder)]
fn main(n: usize, k: usize, a: [usize; n]) -> usize {
    binary_search(1_000_000_000, 0, |m| {
        k >= a.iter().map(|a| (a + m - 1) / m - 1).sum()
    })
}

ということで、これもほぼノイズを消せていると思います。

解決できない問題

input!マクロ一回でデータを読み込むのが面倒なもの（複数のクエリを処理する問題で、クエリタイプごとに構文が変わってくるようなやつとか）は、手動で入力した方が楽だと思います。

インタラクティブな問題も、同じ理由で今のところ対応できていません。ただ、これに関しては何かうまい抽象化の方法がありそうなので、もう少し考えてみたいところではあります。