Leanで競技プログラミングの入力をスッキリ記述するマクロ+α

Last updated at 2026-05-03Posted at 2026-05-03

※ この記事は @tanakh (Hideyuki Tanaka) さんのRustで競技プログラミングの入力をスッキリ記述するマクロにインスパイアされています

2025年10月から AtCoder で Lean v4.22.0 が使えるようになりました。

しかし、Leanで標準入力から望みのデータをパースするのは結構面倒くさいのです。(たぶん)¹

ということで、その辺を改善するマクロを作りました。
↓こちらをコピペして使ってください。
https://gist.github.com/spinylobster/fc70b3aa14cceaea1351ec5e033281a0

practice コンテストのA問題の場合だと、僕のマクロではこう書けます。

/- 省略 (マクロの定義) -/

def main : IO Unit := do
  -- 入力をパースするマクロ
  input!
    a b c : Nat,
    s : String

  -- 結果の出力
  let stdout ← IO.getStdout
  stdout.putStrLn s!"{a + b + c} {s}"

-- 入出力例1のテスト
#test
  input "
1
2 3
test
"
  output "
6 test
"

-- 入出力例2のテスト
#test
  input "
72
128 256
myonmyon
"
  output "
456 myonmyon
"

入出力例を使ったテストも簡単に書けるので、簡易的な正誤チェックに使えます。

このmain関数をマクロなしで書くと大体こんな感じになります。

def main : IO Unit := do
  -- 入力の受け取り
  let stdin ← IO.getStdin
  let input ← stdin.readToEnd
  let words := input.trimRight.split Char.isWhitespace
    |>.toArray

  -- 入力のパース
  let a := words[0]!.toNat!
  let b := words[1]!.toNat!
  let c := words[2]!.toNat!
  let s := words[3]!

  -- 結果の出力
  let stdout ← IO.getStdout
  stdout.putStrLn s!"{a + b + c} {s}"

諸々の記述が省略できるので問題を解くことに集中できて便利なはず。

`input!`マクロの使い方

構文は元記事のRustのマクロがいい感じだったのでそれに寄せています。
実行効率に関しては僕の知識が足りていないので最適化は不十分でしょうが、致命的なほど遅くはないと思っています。
(もっといいやり方を知っている人は、教えてくれると嬉しいです)

まず最初に挙げたような、自然数二つを読み込んで変数にバインドするような例であれば、次のように書けます。

def main := do
  input!
    n m : Nat

  -- solver code...

競技プログラミングによくありがちな、長さnの配列を読み込ませるものであれば、次のように書けます。

-- n
-- v_1 v_2 ... v_n

-- こういうやつを読み込むコード
input!
  n : Nat,
  v : Vector Int n

vはInt型の値をn個含んだ配列になります。
長さ指定のない配列Arrayを使った場合、入力の最後まで読み込むことになります。

配列はネストできるので、n×m行列であれば、

input!
  n m : Nat,
  mat: Vector (Vector Int m) n

こうなります。

ちょっと面倒な文字列をArray Charとして読み込むのも、

input!
  cs : Chars.A -- トークンを Array Char として読む

このように書けます。 (リスト版のChars.LやVector版のChars.V nもあります)
普通にArray Charと書いた時とは挙動が変わるので注意してください。
Array Charの場合は1文字ずつ空白で区切られていないとパースに失敗します。

よくある、二次元テキストとして盤面が与えられるようなものなら、次のように書けます。

-- 3 4
-- ####
-- #..#
-- .#.#

-- こういうやつを Array (Array Char) として読み込むコード
input!
  _h _w : Nat,
  board : Array Chars.A

AtCoderなどでデータを入力していて地味に面倒な、インデックスが1-originとして与えられるときに0-originへの変換なども、自動でできるようにしてあります。

-- n m
-- f_1 t_1
-- f_2 t_2
-- ...
-- f_m t_m

-- こういうグラフを読み込むコード
input!
  n m : Nat, -- n: ノード数, m: 枝数
  edges: Vector (Nat1 × Nat1) m -- Nat1 は 1-origin to 0-origin変換を行う

独自の型へパースしたい場合はForceParse型クラスのインスタンスを作れば対応できます。
デフォルトでは以下の型のインスタンスを作ってあります。

文字・文字列: String, Char
自然数: Nat, UInt8, UInt16, UInt32, UInt64
整数: Int, Int8, Int16, Int32, Int64, ISize
浮動小数: Float
真偽値: Bool
その他: Array α, List α, Vector α n, α × β

注意事項

パースに失敗した時はpanic!しています。
Leanの仕様上、panic!するとパース結果としてはデフォルト値が使われて処理は継続していくのでこれはちょっとわかりにくい挙動だと思います。
気力が湧いたらIOで例外を投げる形に直したいと思います。

`#test`コマンドの使い方

#evalコマンドのようにコード内でmain関数の実行結果を確認するコマンドも作りました。

#test
  input "
1
2 3
test
"
  output "
6 test
"

入力文字列は先頭と末尾にある空白や改行文字をトリミングし、末尾に改行文字を足した上でmain関数に渡されます。
AtCoderの場合も入力の末尾は改行文字らしいのでこれで正しいはず。

出力文字列とmain関数の出力は完全一致で比較します。
こちらも同様に、トリミングした上で末尾に改行文字を足したものと比較されます。

役に立つかわかりませんが、テスト対象にmain関数以外を指定することもできます。

def main' : IO Unit := do
  input!
    _n _m : Nat,
    cs : Array Chars.A

  let stdout ← IO.getStdout

  let result := cs.flatMap id |>.count '.'
  stdout.putStrLn s!"{result}"

#test
  input "
3 4
####
#..#
.##.
"
  output "4"
  main := main'

好きに使ってください

AtCoderでLean使用者が増えることを望んでいるので、今回書いたコードはCC0で公開します。つまり、著作権を主張しません。
そのまま使うのも、改造して使うのも自由です。報告やコピーライトの記載も不要です。

ゆる募: 改善提案

「こうしたらもっと使いやすい or 速くなると思う！」みたいな提案があったら歓迎します。
具体的にどの問題で役立つか、ユースケース付きで教えてくれるとありがたいです。

やりのこしたこと・愚痴

改善できそうなことは色々あるのですが、AtCoderが採用しているLeanのバージョンが新しくなったら取り組もうと思います。
その頃にはきっと何もかも忘れているので、issue代わりにここへ書き残しておきます。

v4.22.0の構文拡張の様子がおかしくて諦めたこと
- 変数宣言用のシンタックスカテゴリdecの定義が変になってる
  - 複数の変数名なのでident+にしたいけど、そうしたら怒られるのでterm にした。意味不明
    - 変数の連なりが関数適用と解釈されるので、そこからidentを取り出す処理を書いている。悔しい
- 元ネタにあるパース元の指定 (source := $s っていう構文) を追加したかった
  - syntax (name := inputMacro) "input!" ("(" "source" ":=" term ")")? colGe dec,* : doElem
  - ↑こんな定義にしたかった
Do構文を組み立てるのにはLean.Elab.Term.Do.Codeが役立ちそうに見えたが、使い方が分からないので断念した
- Do構文のelaborationの仕組みは今では大きく変わってるっぽいのでこれはもう忘れた方が良いかも
String.splitの結果がListで、それをArrayに変換してるので効率悪そう
- 最新版だとStringの仕様が変わってるのでこういうのはなくなるかも
wordsをSubarrayにして一個パースする度に再代入してる
panic!しても処理は続く。IOモナドの中なんだから異常終了した方が良さそう
独自の型からのderiving機構とかも頑張ればできそうな気はした
- 例えば列挙型に対応する文字をいい感じに割り当ててForceParse型クラスを導出するとか
#testコマンドは plausible と被るので変えた方が良いかも？
- とはいえ、プロパティベーステストを書きたくなることってあるのかな？とも思う
実際に使ってみて分かったこと
- 恐らくimport Leanしているせいで70MiBくらいメモリを使ってしまっている
  - 何もimportしてないと9MiBくらい
- 実行時間が、速い人と比べると10倍くらい違う

実を言うとまだ他の問題を全く解いていないので、面倒さを実感していません。 ↩

You get articles that match your needs
You can efficiently read back useful information
You can use dark theme

What you can do with signing up