AtCoderの入力処理コードを自動生成してみた

はじめに

競技プログラミングに参加していて解法が思いつかない時、「とりあえず入力を受け取るコードを書いて落ち着こう」となったことはありませんか。
どんなに難しい問題でも入力を受け取る部分を書くところは何も考えずに書けますからね。

しかし入力部分を何も考えずに手を動かすだけで書けるのであれば、なぜ自動化しないのでしょうか。

ということで実際にやってみました。
autotaker/kyopro

Install

haskell製ツールなのでinstallにはstackもしくはcabalが必要です。

$ git clone https://github.com/autotaker/kyopro.git
$ cd kyopro; stack install

またスクレイピングライブラリ等依存関係の多いライブラリを使っているために初回のビルドに非常に時間がかかります。ご了承ください。

Demo

まず作ったツールの使い方を紹介します。現状はAtCoderにのみ対応しています。
例としてABC 121に対して動かしてみます。

$ mkdir abc121; cd abc121
$ kyopro --generate
[Info] "render.yaml" is generated.
$ kyopro abc121
[Info] "tasks.json" doesn't exists. Trying to scrape it from the contest page.
[Info] Scraped "tasks.json".
[Info] creating directory: A
[Info] creating directory: B
[Info] creating directory: C
[Info] creating directory: D

このようにatcoderのサイトをスクレイピングして問題の数だけディレクトリを作ってくれます。

そしてA/main.cpp, B/main.cpp, C/main.cpp, D/main.cppにそれそれ生成されたコードが出力されます。

A問題

例えばA問題ならこんな感じのコードを生成します。

A/main.cpp

#include<bits/stdc++.h>
void solve(int H, int W, int h, int w) {

}
int main(void) {
  int H;
  int W;
  int h;
  int w;
  std::cin >> H;
  std::cin >> W;
  std::cin >> h;
  std::cin >> w;
  solve(H,W,h,w);
  return 0;
}

B問題

B問題は2次元配列を読み取る必要があります。

B/main.cpp

#include<bits/stdc++.h>
void solve(int C, int M, int N, std::vector<std::vector<int>> A, std::vector<int> B) {

}
int main(void) {
  int C;
  int M;
  int N;
  std::cin >> N;
  std::cin >> M;
  std::cin >> C;
  std::vector<std::vector<int>> A(N, std::vector<int>(M));
  std::vector<int> B(M);
  for(auto i0 = 0; i0 < M; ++i0)
  {
    std::cin >> B[i0];
  }
  for(auto i1 = 0; i1 < N; ++i1)
  {
    for(auto i0 = 0; i0 < M; ++i0)
    {
      std::cin >> A[i1][i0];
    }
  }
  solve(C,M,N,A,B);
  return 0;
}

C問題

今度は2個の整数が$N$個縦に連なっているパターン。

C/main.cpp

#include<bits/stdc++.h>
void solve(int M, int N, std::vector<int> A, std::vector<int> B) {

}
int main(void) {
  int M;
  int N;
  std::cin >> N;
  std::cin >> M;
  std::vector<int> A(N);
  std::vector<int> B(N);
  for(auto i0 = 0; i0 < N; ++i0)
  {
    std::cin >> A[i0];
    std::cin >> B[i0];
  }
  solve(M,N,A,B);
  return 0;
}

D問題

３２ビットに収まらないことに注意する必要があります。

D/main.cpp

#include<bits/stdc++.h>
void solve(long long A, long long B) {

}
int main(void) {
  long long A;
  long long B;
  std::cin >> A;
  std::cin >> B;
  solve(A,B);
  return 0;
}

C++ユーザじゃないんだけど？

はい、私もそうです。
render.yamlがコード生成に使われるテンプレート集なので別の言語への移植も簡単です。

$ kyopro abc121 --render render-haskell.yaml
[Info] "tasks.json" doesn't exists. Trying to scrape it from the contest page.
[Info] Scraped "tasks.json".
[Info] creating directory: A
[Info] creating directory: B
[Info] creating directory: C
[Info] creating directory: D

とするとhaskell用のテンプレートで生成されます。

A/Main.hs

{-# LANGUAGE TypeFamilies, BangPatterns #-}
import Data.Array
import Data.Array.Unboxed
import Data.Array.IO
import Data.Array.Unsafe
import Data.Int
import Control.Monad.State.Strict
import Data.ByteString.Char8 as B
import Data.ByteString(ByteString)
import Data.Char

type Scanner a = StateT ByteString IO a
type family Freeze a where
  Freeze (IOUArray i v) = UArray i v
  Freeze (IOArray i v) = Array i v

scanInt :: Scanner Int
scanInt = state $ \s ->
    let Just (i, s') = B.readInt s in
    (fromIntegral i, B.dropWhile isSpace s')

scanString :: Scanner B.ByteString
scanString = state $ \s ->
    let (s1, s2) = B.span (not . isSpace) s in
    (s1, B.dropWhile isSpace s2)

scanReal :: Scanner Double
scanReal = state $ \s ->
    let Just (d, s') = pure (v, B.dropWhile isSpace s2)
            where
            (s1, s2) = B.span (not . isSpace) s 
            !v = read (B.unpack s1)
    in (d, s')

main :: IO ()
main = B.getContents >>= \content -> flip evalStateT content $ do
    vH <- scanInt
    vW <- scanInt
    vh <- scanInt
    vw <- scanInt
    liftIO $ solve vH vW vh vw
solve :: Int -> Int -> Int -> Int -> IO ()
solve vH vW vh vw = undefined

テンプレートファイルrender.yamlを覗いてみましょう。
Gingerというjinja2方言のテンプレート言語で記述されています。

詳しく説明はしませんがなんとなく雰囲気は伝わると思います。
kyopro --generateで生成される実際のrender.yamlには簡単なドキュメントが付いているので自分で書き換える時には参考にしてください。

render.yaml

tabstop: 2
file: main.cpp
type:
    int32: int
    int64: long long
    real: double
    string: std::string
    unknown: unknown
    array: |
        {% if length(dim) == 1 %}
        std::vector<{{elem}}>
        {% elif length(dim) == 2 %}
        std::vector<std::vector<{{elem}}>>
        {% elif length(dim) == 3 %}
        std::vector<std::vector<std::vector<{{elem}}>>>
        {% endif %}
term:
    var: |
        {{ var }}
    param: i{{ param }}
    decl: |
        {% if length(dim) == 0 %}
        {{ type }} {{ var }};
        {% elif length(dim) == 1 %}
        {{ type }} {{ var }}({{ dim[0] }});
        {% elif length(dim) == 2 %}
        {{ type }} {{ var }}({{  dim[0] }}, std::vector<{{ elem }}>({{ dim[1] }}));
        {% elif length(dim) == 3 %}
        {{ type }} {{ var }}({{  dim[0] }}, std::vector<std::vector<{{ elem }}>>({{ dim[1] }}, std::vector<{{ elem }}>({{ dim[2] }})));
        {% endif %}
    scan: |
        std::cin >> {{ var }}{% for i in index %}[{{ i }}]{% endfor %};
    main: |
        #include<bits/stdc++.h>
        void solve({{vars[0][1]}} {{vars[0][0]}}{% for var in vars[1:] %}, {{ var[1] }} {{ var[0] }}{% endfor %}) {

        }
        int main(void) {
          {{ main }}
          solve({{vars[0][0]}}{% for v in vars[1:] %},{{v[0]}}{% endfor %});
          return 0;
        }
    for: |
        for(auto {{ param }} = 0; {{ param }} < {{ end }}; ++{{ param }})
        {
          {{ body }}
        }

技術的詳細

このツールを実装する際に一番面白いところは入力仕様を頑張ってパーズするところなのでどう実装されているか簡単に説明します。

構文解析

まず、AtCoderの入力仕様はこんな感じのデータとして取得できます。

<var>M</var> <var>N</var>
<var>A_{11}</var> <var>A_{12}</var> <var>...</var> <var>A_{1N}</var>
<var>A_{21}</var> <var>A_{22}</var> <var>...</var> <var>A_{2N}</var>
<var>\vdots</var>
<var>A_{M1}</var> <var>A_{M2}</var> <var>...</var> <var>A_{MN}</var>

これを眺めていると、なんとなくA_{11} A_{12} ... A_{1N}というパターンをパーズできれば良いような気がします。

このパターンを見つけたらSequence { from = 1, to = "N", param = "i", pattern="A_{1i}"}というデータに書き換えてみましょう。

すると次のようなリストに変換されます。

[ Simple ["M", "N"]
, Sequence { from = 1, to = "N", param = "i", pattern="A_{1i}"}
, Sequence { from = 1, to = "N", param = "i", pattern="A_{2i}"}
, VDots
, Sequence { from = 1, to = "N", param = "i", pattern="A_{Mi}"}
]

今度はA_{1i} A_{2i} ... A_{Mi}というパターンが見えてきます。
これを同じように書き換えると

[ Simple ["M", "N"]
, Sequence { from = 1, to = "M", param = "j", 
    pattern = Sequence{ from = 1, to = "N", pattern="A_{ji}" }
  } 
]

みたいなデータに書き換わります。

この処理のアイデアを説明するのは割と簡単だと思いますが、実装しようとすると思いの外面倒なことに気づくと思います。
一般にパターン$P$に出現する変数$i$を式$e$で置き換えたパターンのことを$P[e/i]$と書くことにすると、やりたいことは、「３つのパターン$A$, $B$, $C$から、$A=P[1/i]$, $B=P[2/i]$, $C=P[N/i]$, となるような一般化パターン$P$と変数$N$の組を求める」ということになります。愚直にやると、ありうる一般化パターン$P$の候補を列挙して$A$, $B$, $C$とマッチングさせる必要があり、実装が煩雑になりがちです。

今回はこの部分でHaskellの双方向データ操作ライブラリlensを使っています。このライブラリを使うと、jQueryのセレクタみたいな感覚でデータ構造から欲しい要素を抜き出したり、ピンポイントで書き換えるということができます。今回の実装では、一般化パターンを列挙する代わりに、パターン$A$から$A$に出現する添え字の相対位置を表すTraversalと呼ばれるデータを列挙し、列挙したTraversal達を$B$, $C$に適用することで、マッチングを計算しています。

次元推定

次に、コード生成を行うためには各変数の型が何なのか、特に配列型ならば、その次元とサイズはどれだけかを推論してやる必要があります。
変数の次元は入力仕様の数式の添え字の数を見れば簡単にわかります。例えば$A_{MN}$と書かれていれば次元は２であることが明らかですし、そのサイズが$ M \times N$であることも簡単にわかります。

難しいのは配列の要素の型を推論することです。
問題文には自然言語で型が書かれていますが、私は自然言語処理は素人なのでこれを抽出するのは実装困難です。そこでサンプル入力を当てはめて型を推論するという手法を用いています。

先ほど構文解析したパターンにサンプル入力を当てはめてみましょう。

[ Simple ["M", "N"]
, Sequence { from = 1, to = "M", param = "j", 
    pattern = Sequence{ from = 1, to = "N", pattern="A_{ji}" }
  }
]

sample.txt

3 2
30 20 10
0  -10 3

当てはめると次のような割り当てになります。

{ "M": 3, "N": 2 
, "A_{11}" : 30, "A_{12}" : 20, "A_{13}": 10
, "A_{21}" : 0,  "A_{22}" : -10, "A_{23}": 3
} 

ここから、各値が属する最小の型を計算して当てはめます。

{ "M": "int32", "N": "int32" 
, "A_{11}" : "int32", "A_{12}" : "int", "A_{13}": "int32"
, "A_{21}" : "int32",  "A_{22}" : "int", "A_{23}": "int32"
} 

そして、配列変数の要素の型はこの割り当てられた型のうち最小のものを推論します。

[ Simple [("M",Int), ("N", Int)]
, Sequence { from = 1, to = "M", param = "j", 
    pattern = Sequence{ from = 1, to = "N", pattern=("A_{ji}", Int) }
  }
]

今回の場合は全てint32型になっているので自明に見えますが、

例えば

{ "X_1": "int", "X_2": "string" } 

のような割り当てが見つかった場合、Xの要素の方はint型とstring型のユニオンをとってstringと推論されます。

宣言挿入

最後に配列の確保を行う文を出力するために、各変数がどの時点で有効になるのかを計算してやる必要があります。

[ Simple [("M",Int), ("N", Int)]
, Sequence { from = 1, to = "M", param = "j", 
    pattern = Sequence{ from = 1, to = "N", pattern=("A_{ji}", Int) }
  }
]

このケースでは２次元配列AのサイズはMとNを読み込むまで決定できません。
従って、1行目と２行目の間にAの宣言を挿入します。結果次のようになります。

[ Declare ("M", Int), Declare ("N", Int)
, Simple [("M",Int), ("N", Int)]
, Declare [("A", Array ["M","N"] Int)]
, Sequence { from = 1, to = "M", param = "j", 
    pattern = Sequence{ from = 1, to = "N", pattern=("A_{ji}", Int) }
  }
]

この形までくればコード生成はそれほど難しくはありません。
Declareを変数宣言に、Simpleをリストの長さだけの入力をパーズする文に、Sequenceをfor文に出力すれば良いです。

精度

さて、入力部分を自動生成してくれるツールにはkyuridenamida/atcoder-tooldsという偉大な先人がいます。直近五回のAGC/ARC/ABCでコード生成に成功する率を比較してみました。

Contest Type	atcoder-tools	kyopro
AGC	28/30	25/30
ARC	19/20	19/20
ABC	19/20	19/20

結果、AGCではまだ改善の余地がありますが、ARC/ABCでは同じく成功率95%という高い性能が出ています。

Limitations

最後に現在の実装ではサポートできていない入力パターンをいくつか紹介します。

木の入力

N
a_1 b_1
a_2 b_2
\vdots
a_{N-1} b_{N-1}

このパターンでは列の最後のパターンがN-1となっていて、変数ではないのでうまくいきません。これは実装を多少直せばなんとかなりそうなケース。

配列の配列

N
K_1 A_{11} A_{12} ... A_{1K_1}
K_2 A_{21} A_{22} ... A_{2K_2}
\vdots
K_N A_{N1} A_{N2} ... A_{NK_2}

このケースでは変数Aの見た目上の次元は2ですが、実際には１次元配列の配列として実装する必要があります。これも頑張ればサポートすることは可能だと考えられます。

文脈依存なパターン

N
a_1 a_2 ... a_{2N}

このパターンはなかなかに厄介です。なぜかというとa_{2N}の添え字が、2*Nなのか、２次元配列のa[2][N]を意味するのかの曖昧性があるからです。

まとめ

拙作のAtCoder入力処理コード生成ツールkyoproの紹介を行いました。
まだ荒削りのツールですが、ARC/ABCでは現状でも実用にたる精度が達成できていると思うので興味のある方はぜひ試してみてissueやpull requestを送っていただけると幸いです。

余談

これは全く本質ではないところの苦労なんですが、実装してみると列を表す...が問題によって\dotsだったり\cdotsだったり..だったり、<var></var>で囲まれていたりいなかったりして辛い気持ちになりました。縦向きの...も\vdotsだったり:だったり。（:は縦の省略記号ではないと思うんですが）
作問者の方は少しばかり配慮をしてもらえると助かるのですが。