Zenn
💎

Ruby で競技プログラミング (AtCoder) をやっているときあるある

2022/01/31に公開
Ruby
AtCoder
tech

はじめに

筆者はまだ競技プログラミング歴が浅い初心者である。Ruby が好きで、ふだん仕事でもプライベートでも Ruby を書く機会が多いため、Ruby で競技プログラミングをやっているのだが、競技プログラミングで書くコードはアプリケーションで書くコードとはかなり毛色が異なるため、アプリケーションの実装のノリで計算量が現実的ではないコードを書いてしまったり、Ruby っぽくない書き方をしてしまうことがある。

今回はそんな「Ruby で競技プログラミングをやっているときあるある」について紹介しようと思う。あまり役に立つような記事ではないため、娯楽程度に見てほしい。

用語

  • TLE
    • Time Limit Exceeded (実行時間制限超過) の略
    • 提出したコードが規定の処理時間をオーバーした際に表示される
    • この場合、正解とはならず得点は得られない

Ruby っぽくないループの書き方をしてしまう

Ruby で配列のループを回すときは Array#each を用いるか、Array#map のような関数型言語のようなメソッドを用いることがほとんどだと思う。

そしてこの Array#each を使ったループの回し方は、他の言語にはあまりない[1]ユニークな手法なのではないかと個人的には思っている。他の言語では、for 文など、添字をインクリメントして添字で要素を指定するループの回し方が一般的だと思う。

ところが、競技プログラミングをやっている最中に限っては、トリッキーな位置の要素を指定することが多いため、Array#each のような、順番に要素のみを参照するメソッドではなく for 文や Integer#upto などを利用する頻度のほうが高かったりする。

たとえば、AtCoder Beginner Contest 237 の C 問題の解説 (C++) の 30 〜 35 行目[2]を見てみる。都合によりインデントを一段回上げている。

for (int i = x; i < (n - y); i++) {
	if (a[i] != a[x + n - y - i - 1]) {
		cout << "No" << endl;
		return 0;
	}
}

ループの開始が x で、終了が (n - y) というトリッキーな指定となっている。さらにその中の if の条件式で、a[x + n - y - i - 1] という、かなり複雑な添字の指定になっている。

これと同等のことを Array#each を使って表現するのはなかなか難しいだろう。もちろん添字を使うときのために Array#each_with_index というメソッドも存在するが、結局、添字しか使わないし、ループの開始と終了が 0 〜 N のように単純ではないので、ここでは筆者は以下のように Integer#upto を使いたくなるところだ。

x.upto(n - y - 1) do |i|
  if word[i] != word[x + n - y - i - 1]
    puts 'No'
    exit
  end
end

ちなみにこの項目を書くにあたって、改めて Ruby のループの書き方についておさらいしたのだが、Ruby にはこんなにループの書き方があるのかと驚かされた。

【Ruby入門】ループ処理まとめ（for・times・while・each・upto・downto・step・loop）

業務で使用するのはやはり専ら Array#each くらいだ。

Range#each で代用できてしまうことが判明

記事を書き終わったあとに RuboCop にかけてみて気づいたのだが、添字をレシーバにして Range#each を使用すれば意外にこれだけでもいけてしまうということがわかった。

たとえば、上記の Integer#upto を用いたコードを Range#each に書き換えると以下のようになる。

(x..(n - y - 1)).each do |i|
  if word[i] != word[x + n - y - i - 1]
    puts 'No'
    exit
  end
end

他にも、for を使った以下のような書き方も Range#each が使えそうだ。

arr = %w[apple banana orange]
n = arr.length

# for
for i in 0...n
  puts arr[i]
end

# Range#each
(0...n).each do |i|
  puts arr[i]
end

for を Range#each に書き換えるのはまだ自然だが、Integer#upto を Range#each に書き換えるとカッコの数が増えて若干複雑になるから個人的には Integer#upto を使うかな。

計算量を気にせずに便利なメソッドを使ってしまう

Ruby には (特に文字列や配列において) 他の言語にはないような便利なメソッドが数多くある。ループ内で多用しない限りは便利に使えるメソッドだったりするのだが、こと競技プログラミングにおいては計算量 (処理時間) が仇となって TLE になってしまうことが多発する。筆者も問題を解いているとき、かなりの頻度で TLE が発生するので、普段のコーディングで如何にパフォーマンスを軽んじているかを痛感してしまう。

たとえば、AtCoder Beginner Contest 237 の D 問題。与えられた文字列の文字種に応じて数値の挿入処理を行うというものだ。

配列の特定の位置に要素を挿入するのだから、Rubyist の素直な脳で考えると、Array#insert が使えそうだなとなる。実際に解いたコードがこちら。

#!/usr/bin/env ruby

def main
  n = gets.to_i
  s = gets.chomp
  a = [0]

  pos = 0
  for i in 0...n
    pos += 1 if s[i] == 'R'
    a.insert(pos, i + 1)
  end

  puts a.join(' ')
end

main

これ以上なく簡潔に書けたのではないかと脳内で自画自賛しつつこのコードを提出したところ、TLE となった。原因は a.insert(pos, i + 1) の部分だ。

この Array#insert と同等の処理を行うメソッドを自作しようと思うとわかるのだが、指定した要素を挿入する位置よりも後ろにある要素の位置をすべて 1 つずつずらした上で挿入する必要があるので、それなりに時間がかかる。最悪の計算量は O(n) になると思う。それをループの中で何回も実行しているのだから、最悪の計算量は O(n^2) となり、これでは TLE になっても無理はない。

Array#insert のような便利メソッドがない言語なら、計算量的に時間がかかりすぎるということにすぐ気付けるため、そもそもこれと同等の処理を行うアルゴリズムを実装しようとは思わないが、Ruby だとこれがすぐに書けてしまうので、ついうっかり使ってしまうのである。

さて、一応、計算量的にも問題ないアルゴリズムも紹介する。解法はいくつか存在するようだが、上記のコードに近い実装の一つとして、2 つの空配列を用意して、挿入位置の左側と右側に分けてそれぞれ要素を追加していく方法がある。このアルゴリズムの詳細な解説は 公式サイト に上がっているのでこちらも参照してほしい。

このアルゴリズムで実装したコードがこちら。

#!/usr/bin/env ruby

def main
  n = gets.to_i
  s = gets.chomp
  l = []
  r = []

  for i in 0...n
    if s[i] == 'L'
      r.unshift(i)
    else
      l.push(i)
    end
  end

  puts "#{l.join(' ')} #{n} #{r.join(' ')}".strip
end

main

コード的には記述量が増えて最初に書いたコードほどスッキリはしていないが、こちらのコードであれば TLE にならずに解を求めることができる。

記述量が増えてスッキリしないと書いたが、それはそもそも Array#insert の中身を自分で書いていないからであって、もしこのメソッドのアルゴリズムも自前で書くことになったらきっとそのほうが複雑になるだろう。でも、そうとは感じさせないほど気軽に使えてしまうあたりが Ruby らしくもある。

また、この問題を解いたことにより、Array#insert は計算時間がかかるので、Array#unshift や Array#push が使えるならそちらを使おうという教訓が得られた。

便利なメソッドに頼りすぎて勉強にならない

これは先ほどの「計算量を気にせずに便利なメソッドを使ってしまう」と対比している。計算量的にも問題なく、かつシンプルに書けてしまうのだが、その書き方では競技プログラミングをしている意味がない気がする……、というようなことがまあまあある。

たとえば、AtCoder Beginner Contest 237 の B 問題。与えられた行列の転置行列を求めるというものだ。

他の言語では 2 重ループで愚直に要素を入れ替えて転置行列を求める[3]が、Ruby だと Array#transpose というドンピシャなメソッドが存在する。これを使うと、以下のように書けてしまう。

#!/usr/bin/env ruby

def main
  h, _w = gets.split.map(&:to_i)
  a = []
  h.times do
    a << gets.split.map(&:to_i)
  end

  a = a.transpose

  a.each do |b|
    puts b.join(' ')
  end
end

main

main メソッドの上のほうは入力値を変数に代入する処理で、下のほうは結果を出力する処理である。つまり、この問題の実質的なコードは a = a.transpose のたった 1 行のみとなってしまう。

このコードはもちろん正解だし、TLE にもならない。しかし、競技プログラミングをやって数学的な脳を養う、またはアルゴリズムを勉強するという観点でいうと本当にこれで良いんだろうか (いや良くない) という気持ちになる。

さいごに

今後も Ruby で競技プログラミングを続けていこうと思っているので、また何か気づいたあるあるがあれば加筆する。

脚注

  1. あまりないと書いた直後、JavaScript の for ... of を思い出してしまったので、そうでもないのかな？ とも思ってしまった。それでも、競技プログラミングで良く使われる C++ や Python などでは、添字をインクリメントして添字で要素を指定することが多い気がする。 ↩︎

  2. 2022 年 1 月 31 日現在。 ↩︎

  3. 他の言語にも転置行列を求めるメソッドやライブラリはあると思うが未調査。 ↩︎

GitHubで編集を提案

Discussion