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

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

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

今日はここまで！

いものやま。

雑多な知識の寄せ集め

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

問題

自分の解答