Rubyの繰り返し処理が思ったより遅かった
競プロの演習問題で多重ループで全探索する問題があったので Ruby で実行したらTLEしまくった。
最近発売になった「問題解決のための『アルゴリズム×数学』が基礎からしっかり身につく本」を読んでいる。競プロは C++ でやるのがセオリーなのは百も承知なのだけど、個人的な書きやすさの好みで Ruby を使っている。
件の問題は割とシンプルなもので、N 枚のカードから5枚を選ぶときに、カードに書かれた整数の和が 1000 になるものが何通りかを答える問題。N の制約が 5 <= N <= 100 なので、5秒以内に全探索可(著者環境だと N=100 でも 0.087 秒で終わったそう)らしいのだけど、Ruby だと普通にTLEして泣いた。
Ruby で書いたコード
N = gets.to_i
A = gets.split.map(&:to_i)
cnt = 0
(0...N).each do |i|
((i+1)...N).each do |j|
((j+1)...N).each do |k|
((k+1)...N).each do |l|
((l+1)...N).each do |m|
cnt += 1 if A[i] + A[j] + A[k] + A[l] + A[m] == 1000
end
end
end
end
end
puts cnt
Range使ったeachが悪いのかと思ってwhileで書き直してみたけどいくつかの large パターンのケースで5秒を超えてしまう。
試しに Crystal で同じ感じで書いてみたら実行時間が 1/10 くらいになって余裕で通った。
Crystal で書き直した版
N = read_line.to_i
A = read_line.split.map(&.to_i)
cnt = 0
(0...N).each do |i|
((i+1)...N).each do |j|
((j+1)...N).each do |k|
((k+1)...N).each do |l|
((l+1)...N).each do |m|
cnt += 1 if A[i] + A[j] + A[k] + A[l] + A[m] == 1000
end
end
end
end
end
puts cnt
実際のコンテストで頻繁に遭遇するケースではないと思うけど、Ruby で競プロやろうとすると壁があるというのは本当らしい。万が一に備えて Crystal で書くことも考えた方がいいのかもしれない。それだったら素直に C++ 使えという話ではあるが・・・