Rubyで競技プログラミング
この記事はRuby Advent Calendar 2016の最終日の記事です。
発端
秋に行われたRubyKaigi 2016のあとのAfter partyで、以下のような発表をしました。
内容は、会社の同僚と一緒にRubyでAtCoderをやっているという話でした。AtCoderはオンラインで競技プログラミングができるサイトで、Rubyを含めたさまざまな言語で参加することができ、過去問についても解答を受け付けているため好きなときにチャレンジすることができます。
競技プログラミングではC, C++, C#, Javaといった言語が主流なので、Rubyでは速度が足らないのではないか？と思われるかもしれません。実際、AtCoderでは2秒の制限時間に対し、入力データが10万行を超えることもあります。それでも、適切なアルゴリズムを使うことで、9月時点でチャレンジしていた範囲(ABC035〜041)はRubyでもちゃんとクリア(AC)することができました。
…ただ一問を除いては。
ABC037 D 経路
問題の一問がこれです。数字が敷き詰められた盤面を「数字が大きくなっていくように歩く」とき、何通りの歩き方があるかという問題です。(今回は実行速度の話なので具体的な解法については触れませんが、ネタバレが気になる場合は以下を読む前に一度チャレンジしてみて下さい)
この問題だけがどうしても2秒以内に解けませんでした。TLE(Time Limit Error)だけでなくRE(Runtime Error)が出ていることから、stack level too deepが出ていることが予想されます。アルゴリズムは再帰呼び出しを使ったもので、盤面は最大で1000×1000マスなので、デフォルトのスタックサイズをオーバーしてしまうことは十分に考えられます。
Rubyとスタックサイズ
RubyKaigiでこの件についてささださんに尋ねてみたところ、「スタックサイズは環境変数経由で設定することができる」という情報を頂きました。
ただし残念ながら、環境変数を使わずにスクリプト内からこれを設定する方法はまだないとのことでした(理由は「特に要望がなかったから」)。AtCoderには投稿したプログラムの環境変数をするようなインターフェイスはないので、この方法は使えなさそうです。
と思いきや…
@yhara https://t.co/W5KMdQCyoB 一度スタックサイズを増やしてから再度実行し直すことで無事通すことが出来ました！🎉
— Shuichi Tamayose (@_simanman) 2016年9月14日
冒頭のスライドを見た@_simanmanさんから、クリアできたという報告が！
どうやら、以下のようにKernel#execでRubyインタプリタをもう一度起動させることで、環境変数を指定しているようです。そんなことをしたら2秒制限にひっかかりそうですが、1991msというぎりぎりのタイムでクリアできています。
if !ENV['RUBY_THREAD_VM_STACK_SIZE']
exec({'RUBY_THREAD_VM_STACK_SIZE'=>'1000000000'}, '/usr/bin/ruby', $0)
...
どうすれば良かったのか？
クリアできることが分かったので、自分が投稿したプログラムと比較して、どこが悪かったのかを調べてみました。
1. Bignumを経由する箇所があった
まず単純なミスとして、問題文の「10^9+7で割った余りを求めてください」という一節が実装できていませんでした。このため盤面が大きい場合に間違った答えが出る上に、答えが出るまでにとても時間がかかるようになっていました。
Rubyでは四則演算の結果が大きい数になると、自動的にFixnumからBignumへ結果が拡張されるようになっています。これはとても便利な機能ですが、Bignumの計算はFixnumよりも低速です。今回は結果の総数ではなく10^9+7で割った余りだけ分かれば良いので、計算の途中でこまめに10^9+7で割るようにすれば、Bignumを経由しなくて済みます。
2. スタックを浅くする必要があった
AtCoderでは時間制限の他に、メモリ量の制限もあります。ABC037 D問題では使えるメモリは256MBまでで、これを超えるとMLE(Memory Limit Error)というエラーになります。
いろいろ試してみた結果、筆者のもとのプログラムだとスタック制限を回避できたとしてもMLEで止まってしまうことが分かりました。問題の箇所はここです。
def calc(pos)
ret = 1
DIRS.each do |d|
if MAP[pos] < MAP[pos+d]
plus = ANS[pos+d] || calc(pos+d) # 再帰呼出し
...
このDIRS.eachをwhileに置き換えると、無事Acceptedになります。
def calc(pos)
ret = 1
i = 0
v = MAP[pos]
while i < 4
d = DIRS[i]
if v < MAP[pos+d]
plus = ANS[pos+d] || calc(pos+d)
...
@_simanmanさんのプログラムではwhileの代わりにループアンローリングを使っています。ループアンローリングは主に高速化のためにループをべた書きに展開するという手法で、例えば上のwhileを展開すると以下のようになります。
def calc(pos)
ret = 1
if MAP[pos] < MAP[pos+1]
plus = ANS[pos+1] || calc(pos+1); ret += plus
end
if MAP[pos] < MAP[pos-1]
plus = ANS[pos-1] || calc(pos-1); ret += plus
end
if MAP[pos] < MAP[pos+W]
plus = ANS[pos+W] || calc(pos+W); ret += plus
end
if MAP[pos] < MAP[pos-W]
plus = ANS[pos-W] || calc(pos-W); ret += plus
end
まとめ
今回は以上です。一見どうしようもなさそうな状況でも、意外な方法で回避できるというのが面白かったです。
- スタックサイズは環境変数で増やせる
- 途中の計算がBignumを経由していて遅かった
- eachではなくwhileを使うとスタックが浅くなる
- Rubyでも競技プログラミングはできるぞ！(少なくともAtCoder Beginner Contestの範囲は)
