CodeIQ:撤去作業の果てに現れる数列

私自身が表題の問題を解いた時のプログラムについて解説します。
問題の詳細は「【撤去作業の果てに現れる数列」(CodeIQ)を参照してください。

問題の概要

問題を引用します。
【概要】
0から始まる無限に続く整数の列がスタートです。
その列から「素数の ni 個前を撤去する」という操作を繰り返します。
操作の結果得られる列の、最初の10項をコンマ区切りで出力してください。

【詳細】
{ ni } を表す文字列が3 2 3のように標準入力からやってきます。空白区切りの 10進数 です。
この入力の場合は、下表のように列は変化します:

0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,...
↓ n1=3 なので、0,2,4,8,10,14,16,20,26,28,...などを撤去する。
1,3,5,6,7,9,11,12,13,15,17,18,19,21,22,23,24,25,27,29,...
↓ n2=2 なので、1,5,7,11,13,17,21,25,29,35,37,41,...などを撤去する。
3,6,9,12,15,18,19,22,23,24,27,30,31,32,33,36,39,42,43,46,...
↓ n3=3 なので、12,18,24,36,42,48,54,...などを撤去する。
3,6,9,15,19,22,23,27,30,31,32,33,39,43,46,47,49,52,53,57,...

結果、最初の 10項は
3,6,9,15,19,22,23,27,30,31
となるので、これを出力すればよいというわけです。

【例】
入力出力
3 2 33,6,9,15,19,22,23,27,30,31
2 34,6,7,10,12,13,16,18,20,22

【補足】
不正な入力に対処する必要はありません。
入力に含まれる数は、1〜100 の整数です。
入力に含まれる数字の数は、1〜20個 です。
素数の判定は、ライブラリがあればそれを利用して構いません。

私のプログラム

Rubyで解答しています。

#!/usr/bin/ruby

require 'prime'

$Primes = Hash.new(false)
$Last = 0
$Max = 10000

def init()
	Prime.each{|pr|
		if pr > $Max then break end
		$Primes[pr] = true
		$Last = pr
	}

	return (0..$Last).to_a
end

def solve(aheads)
	results = Array.new(aheads.size+1){[]}
	results[0] = $Base

	for i in 0...aheads.size
		for j in 0...results[i].size
			v = results[i][j]
			results[i+1] << v

			if $Primes[v] then
				r = j-aheads[i]
				if r >= 0 then results[i+1][r] = nil end
			end
		end
		results[i+1].compact!
	end

	return results.last[0...10]
end

# main
$Base = init()
while line = gets
	line.strip!
	if line.empty? then next end

	input = line.split.map{|a| a.to_i}
	puts solve(input).join(",")
end

解説

きちんとやろうと思うと難しいです。
私のコードはきちんとしていません。

考え方

まず、私のコードの何がダメかを述べておきます。
私のコードは10000以下の素数の範囲だけを調べて結果を返しています。なので、もし10000より大きい素数を考慮しなければいけないテストケースがあればNGになってしまいます。要するに、上限を勝手に固定した手抜きプログラムというわけです。

この手抜きプログラムだとコードは非常に簡単で、問題文の通りに実装すれば良いだけです。

main

init()を呼んで必要な準備をし、入力値を数値の配列にしてsolve()に渡します。
solve()が結果を配列で返すのでそれをフォーマットして印字して終わります。

init()

数字が現れるたびに素数判定をするのが嫌なので素数表を作ってしまいます。これはRubyの素数ライブラリを使って列挙し、$Primesに{素数=>true}の形式で保持します。こうしておくことで$Primes[値]がtrueなら素数であることをO(1)で判定できるようになります。
同時に、0〜10000以下の最大の素数の配列を作って返します。これが問題の処理をするための初期データになります。

solve(map)

問題文の通りの処理をします。
引数aheadsは入力値です。
resultsは[0]が0以上、10000以下の最大の素数までの連続した数、[1〜]が問題文の通り1回目〜の除外処理をした結果を保持する領域です。

23〜34行目が除外処理です。
24〜32行目で1回前の数列から素数のX個前の数字を除外する処理をしています。処理手順は一旦、1回前の数字をそのまま結果の配列に追加し(25〜26行目)、追加した数が素数なら(28行目)削除位置の数をnilに置換します(29〜30行目)。
最後まで終わってからArray#compact!でnilを削除します(33行目)。
最終的にresultsの最終要素の前から10要素を返せば答えになります。

一旦nilにしてからまとめて削除するのはちょっとした工夫です。Rubyの配列の実装の内容は知りませんが、配列の要素削除、特に前方の要素削除は2つの意味で厄介です。
まず、削除が発生するたびに削除位置より後ろの要素を前方にコピーする処理が発生する可能性があります(C++のvectorはこうなります)。こうなると処理速度が極端に遅くなってしまいます。
もう一つは参照位置がずれることです。私のプログラムはRubyらしくなく配列要素番号でのアクセスなのではっきり問題になりますが、Rubyの配列の内部実装がC/C++の配列やvectorを利用しているなら同じ問題を生じます。例えば[0,1,2,3,4,5]という配列があって3(4番目の要素)を参照している時にの2つ前を削除したとします。そうすると[0,2,3,4,5]になりますが、4番目の要素は4にずれてしまいます。このまま、次の要素を取り出したら4を飛ばして5になってしまうわけです。この問題を避けるためには配列を逆順([5,4,3,2,1,0])にして前方を削除するようにすれば良いのですが、正しく処理できないため今回はこの手を使えません。例えば1つ前の数を削除する場合、後ろからやると[5,3,1,0]で昇順に戻すと[0,1,3,5]です。しかし、期待される結果は[0,3,5]です。つまり、後ろからやると2を処理する前に2が削除されるので正しい結果を得られないことになってしまいます。
削除するべき数を一旦nilにすることでこれらの問題を回避できるというわけです。

雑感

さて、よくないプログラムを書いてしまいましたが、きちんとやるにはどうすべきでしょうか。
一番、楽なのは入力から最後の数を計算してしまうことです。なんとなく、誤差で多少大きくなっても良いが小さくなるのはダメ程度で最後の数を予測することはできそうな気はするのです。が、やり方がよくわかりません。

最後の数を計算できないとなると次の様なロジックが思いつきます。例えば入力値が[3 2 4]とします。

3個前の削除ができるところまで数を進めて削除します([0,1,2,3]→[1,2,3]になる)。
この配列から2個前の数を削除できるか試します(この場合は3が素数なので可能で[2,3]になる)。
同じく4個前の数が削除できるか試しますが、今度はできないので最初に戻ります。

次の素数まで増やします([2,3,4,5])。そして同じことを繰り返します。
3個前の数を削除し[3,4,5]、2個前の数を削除し[4,5]、4個前の数はないので戻ります。
数を増やして[4,5,6,7]、3個前の数を削除し[5,6,7]、2個前の数を削除し[6,7]、4個前の数はないので戻ります。
数を増やして[6,7,8,9,10,11]、3個前の数を削除し[6,7,9,10,11]、2個前の数を削除し[6,7,10,11]、4個前の数はないので戻ります。
数を増やして[6,7,10,11,12,13]、3個前の数を削除し[6,7,11,12,13]、2個前の数を削除し[6,7,11,13]、4個前の数はないので戻ります。
数を増やして[6,7,11,13,14,15,16,17]、3個前の数を削除し[6,7,11,13,15,16,17]、2個前の数を削除し[6,7,11,13,16,17]、4個前の数を削除し[6,11,13,16,17]。入力値の分だけ削除処理をしたので戻って数を増やします。
以降、この繰り返しで配列の要素数が10になるまで繰り返します。

最初、この方法でプログラムを考え始めたのですが、考えている間になんとなく問題の条件だと10000は超えない様に思えたので10000以下に決めうちしてました ^^;