読者です 読者をやめる 読者になる 読者になる

いものやま。

雑多な知識の寄せ集め

CodeIQ「ディビジョン・ナイン」問題を解いてみた。

Ruby

CodeIQで出題された「ディビジョン・ナイン」問題。

Rubyで解いてみたので、コードを公開してみる。

問題

問題は以下のとおり:

1から4の数字を使って  n 桁の整数を作ります。
このとき、9の倍数となるものを考えましょう。

例えば  n = 3 であれば、234、333、441、などが9の倍数です。
必ずしも1から4の全ての数字を使う必要はありません。

1から4の数字を使って作る  n 桁の整数のうち、9の倍数となるものの個数を  F(n) と定義します。

例えば、 F(1) = F(2) = 0 F(3) = 10 F(4) = 40 となることが確かめられます。

標準入力から、自然数  n \: (1 \le n \le 20) が与えられます。
標準出力に  F(n) の値を出力するプログラムを書いてください。

自分の解答

9の倍数の判定は簡単で、各桁の数字を足して9で割り切れれば、9の倍数。
ということで、関数の再帰呼び出しを使った深さ優先探索のコードをまず書いた:

def f(n, nums)
  if n == 1
    rest = nums.inject(0, &:+) % 9
    (5 <= rest && rest <= 8) ? 1 : 0
  else
    (1..4).inject(0) do |result, t|
      result += f(n-1, [nums, t].flatten)
    end
  end
end

n = $stdin.gets.to_i
puts f(n, [])

関数  f(n, nums) は、使う数字が決まっていない桁数  n、各桁で使うことにした数字の配列  nums に対して、9で割り切れる数が何個あるかを返すもの。
例えば、 f (2, [1, 1]) であれば、11xx(xはそれぞれ1〜4の数字)という数のうち9で割り切れるものの個数、 f (1, [1, 1, 1]) であれば、111x(xは1〜4の数字)という数のうち9で割り切れるものの個数が返ってくる。

 n が1の場合は、使うことにした数字を全て足して9で割った余りが5〜8の間に入っていれば、9で割り切れる数を作ることができる。
例えば、割った余りが5になっていれば、最後の1つの数字として4を使えば、9で割り切れる数になる。
なので、そういう場合には1を、そうでない場合には0を返している。

そして、 n が2以上の場合は、使う数字を1つ固定して再帰呼び出しすれば、1〜4を使った場合のそれぞれの9で割り切れる数の個数が分かるので、それを合計して返せばいい。

ということで、論理的にはこれでOKなんだけど、実際に実行すると、めちゃくちゃ時間がかかる(^^;
そこで、もうちょっと高速化を。

上のコードでは、各桁で使うことにした数字を配列のまま持っていて、最後 ( n = 1 のとき)に足し算を行っているけど、実際に必要なのは、各桁で使うことにした数字の合計を9で割った余りだけ。
そこで、次のように修正:

def f(n, rest)
  if n == 1
    (5 <= rest && rest <= 9) ? 1 : 0
  else
    (1..4).inject(0) do |result, t|
      result += f(n - 1, (rest + t) % 9)
    end
  end
end

n = $stdin.gets.to_i
puts f(n, 0)

 rest は各桁で使うことにした数字の合計を9で割った余り。
再帰呼び出しをするときには、使うことにした数字の配列を作って渡すんじゃなくて、使うことにした数字を  rest に足して9で割った余りを渡すようにしている。

これで配列を生成するコストは減ったけど、まだまだ遅い。

ただ、ここで  1 \le n \le 20 0 \le rest \le 8 であることを考えると、同じ引数での関数呼び出しが何度も行われている可能性が高い。
となれば、キャッシュを使ってやればかなり高速化するはず。

ということで、キャッシュを使うようにしたコードが以下:

@f_cache = [
  #  0,   1,   2,   3,   4,   5,   6,   7,   8, (rest)
  [nil, nil, nil, nil, nil, nil, nil, nil, nil], # 0
  [  0,   0,   0,   0,   0,   1,   1,   1,   1], # 1
  [nil, nil, nil, nil, nil, nil, nil, nil, nil], # 2
  [nil, nil, nil, nil, nil, nil, nil, nil, nil], # 3
  [nil, nil, nil, nil, nil, nil, nil, nil, nil], # 4
  [nil, nil, nil, nil, nil, nil, nil, nil, nil], # 5
  [nil, nil, nil, nil, nil, nil, nil, nil, nil], # 6
  [nil, nil, nil, nil, nil, nil, nil, nil, nil], # 7
  [nil, nil, nil, nil, nil, nil, nil, nil, nil], # 8
  [nil, nil, nil, nil, nil, nil, nil, nil, nil], # 9
  [nil, nil, nil, nil, nil, nil, nil, nil, nil], # 10
  [nil, nil, nil, nil, nil, nil, nil, nil, nil], # 11
  [nil, nil, nil, nil, nil, nil, nil, nil, nil], # 12
  [nil, nil, nil, nil, nil, nil, nil, nil, nil], # 13
  [nil, nil, nil, nil, nil, nil, nil, nil, nil], # 14
  [nil, nil, nil, nil, nil, nil, nil, nil, nil], # 15
  [nil, nil, nil, nil, nil, nil, nil, nil, nil], # 16
  [nil, nil, nil, nil, nil, nil, nil, nil, nil], # 17
  [nil, nil, nil, nil, nil, nil, nil, nil, nil], # 18
  [nil, nil, nil, nil, nil, nil, nil, nil, nil], # 19
  [nil, nil, nil, nil, nil, nil, nil, nil, nil], # 20
]

def f(n, rest)
  cache = @f_cache[n][rest]
  if cache.nil?
    @f_cache[n][rest] =
      (1..4).inject(0) do |result, t|
        result += f(n - 1, (rest + t) % 9)
      end
  else
    cache
  end
end

n = $stdin.gets.to_i
puts f(n, 0)

関数  f の各引数に対するキャッシュを用意して、まだ求まっていないところにはnilを入れてある。
そして、 n = 1 のときの  f の値は分かってるので、最初から入れてある。
こうすると、 n が1かそうでないかの判定は不要になり、キャッシュに値があればその値を返し、そうでなければ値を求めてキャッシュに値を保存し、その値を返せばいいだけになる。

これで十分に速くなったので、用意されたテストケースもちゃんと時間内に解けるようになった。

ちなみに、これは動的計画法になってるみたい。
確かに、1つの桁の数字を決めてより小さい問題を解いていて(分割統治法)、また、より小さい問題のそれぞれの結果をキャッシュしてる(メモ化)。
動的計画法という名前自体は(ナップザック問題などで)知っていたけど、具体的なアルゴリズムのイメージは掴めてなかったので、自分の書いたコードが振り返ってみれば動的計画法になっていたというのは、ちょっと変な感じ。
強化学習の動的計画法とは、かなり違うしね。

少し補足しておくと、強化学習の動的計画法の方は、状態関数の値を他の状態関数の値を使って求めていて(ブートストラップ)、各状態関数の値を保存している(メモ化)という意味では、(定義次第では)動的計画法の条件を満たしていると言えなくもない。
ただ、関数の各値の関係を有向グラフにしたときにループが存在することがあるので、再帰で解くことが出来るとは限らない。
(つまり、厳密に言うと、より小さい部分問題に分割出来るとは限らない)
そこで、(ループしていない場合も)再帰を使って解くのではなく、ヤコビ法で反復を使って問題を解いている。
なので、ナップザック問題を解く動的計画法と、強化学習の動的計画法だと、かなりイメージが違う感じ。

今日はここまで!