Tree構造を利用して問題を解いてみる(Ruby)

前回実装した Tree構造を使って、プロジェクト・オイラーの問61を解いてみました。問題はこれです。

require 'bundler/setup'
require 'utils'    #自作の野良Gem。下では Float#integer? を使っている
require './tree'

def triangle?(a); ((-1 + (1 + 8 * a) ** 0.5) / 2).integer?; end
def square?(a); (a ** 0.5).integer?; end  
def pentagonal?(a); ((1 + (1 + 24 * a) ** 0.5) / 6).integer?; end
def hexagonal?(a); ((1 + (1 + 8 * a) ** 0.5) / 4).integer?; end
def heptagonal?(a); ((3 + (9 + 40 * a) ** 0.5) / 10).integer?; end
def octagonal?(a); ((1 + (1 + 3 * a) ** 0.5) / 3).integer?; end

class Node
  def initialize(num, type)
    @num = num
    @type = type
  end
  attr_reader :type, :num
  
  def upper;  @num / 100; end
  def downer; @num % 100; end
  
  #特定のtypeの全ノードを配列で返す
  def self.select(n, type1)
    n.select {|node| node.type == type1}    #Array:要素はNode
  end
  
  #nodeのあとにselfが続けばtrue、続かなければfalse
  def succeed?(node)
    node.downer == self.upper
  end
  
  def to_s
    "(#{@num}, #{@type})"
  end
  alias :inspect :to_s
end

class Tree
  #node以上のすべてのノード(:rootは除く)のtypeを配列で返す
  def get_types(node)
    get_nodes(:root, node)[1..-1].map {|node| node.type}
  end
end

#typesの配列arにないtypeを配列で返す
def select_unused_types(ar)
  (0..5).to_a - ar
end

## main ##
n = []
for i in 1000..9999
  n << Node.new(i, 0) if triangle?(i)
  n << Node.new(i, 1) if square?(i)
  n << Node.new(i, 2) if pentagonal?(i)
  n << Node.new(i, 3) if hexagonal?(i)
  n << Node.new(i, 4) if heptagonal?(i)
  n << Node.new(i, 5) if octagonal?(i)
end
m = []    #タイプごとに全ノードが入る配列
6.times {|i| m[i] = Node.select(n, i)}

t = Tree.new(:root)
m[5].each {|node| t.add(:root, node)}
#深さ優先探索
t.dfs(:root) do |node, depth|
  if depth == 6    #最下層か
    ar = t.get_nodes(:root, node)
    if ar[1].succeed?(node)    #完全に循環するか
      p ar[1..-1]
      p ar[1..-1].map {|x| x.num}.inject(:+)
      t.put_out
      exit    #終了
    end
  end
  next if node == :root
  
  select_unused_types(t.get_types(node)).each do |type|    #上の階層で使われていないタイプを選択
    m[type].each do |node1|
      next unless node1.succeed?(node)    #数字の「しりとり」になっているか
      node1 = Node.new(node1.num, node1.type) if t.exist?(node1)
      t.add(node, node1)    #ノードの追加
    end
  end
end

深さ優先探索でトラバースしています。
問題にしてはえらく長いコードに見えますが、実質的なループは最後の 20行くらいです。多くの部分はクラスの設定や初期設定が占めているわけです。


結果は下です。

[(1281, 5), (8128, 3), (2882, 2), (8256, 0), (5625, 1), (2512, 4)]
28684

user      0m0.252s

答えが見つかった時点での Tree は以下です。

-- root -- (1045, 5) -- (4560, 0) -- (6084, 1) -- (8400, 2) 
                     -- (4510, 2) -- (1035, 0) -- (3553, 4) -- (5329, 1) 
                                                            -- (5356, 3) -- (5625, 1) 
                                  -- (1081, 0) -- (8100, 1) 
                                               -- (8128, 3) -- (2809, 1) 
                                                            -- (2839, 4) -- (3969, 1) 
                                  -- (1024, 1) -- (2415, 0) -- (1540, 3) 
                                                            -- (1525, 4) -- (2556, 3) 
                                               -- (2485, 0) 
                                               -- (2415, 3) -- (1540, 0) 
                                                            -- (1596, 0) 
                                                            -- (1525, 4) -- (2556, 0) 
                                  -- (1089, 1) -- (8911, 0) -- (1128, 3) -- (2839, 4) 
                                                            -- (1177, 4) 
                                               -- (8911, 3) -- (1128, 0) -- (2839, 4) 
                                                            -- (1176, 0) 
                                                            -- (1177, 4) -- (7750, 0) 
                                               -- (8910, 4) -- (1035, 0) 
                                                            -- (1081, 0) -- (8128, 3) 
                                                            -- (1035, 3) -- (3570, 0) 
                                  -- (1035, 3) -- (3570, 0) -- (7056, 1) -- (5688, 4) 
                                               -- (3553, 4) -- (5356, 0) -- (5625, 1) 
                                                            -- (5329, 1) -- (2926, 0) 
                                  -- (1071, 4) -- (7140, 0) -- (4096, 1) 
                                                            -- (4005, 3) 
                                               -- (7140, 3) -- (4005, 0) 
                                                            -- (4095, 0) 
                                                            -- (4096, 1) 
                     -- (4560, 3) -- (6084, 1) -- (8400, 2) 
                     -- (4558, 4) -- (5886, 0) -- (8649, 1) -- (4950, 3) -- (5017, 2) 
                                               -- (8626, 2) -- (2601, 1) 
                                               -- (8646, 3) -- (4624, 1) 
                                                            -- (4676, 2) 
        -- (1160, 5) -- (6084, 1) -- (8400, 2) 
        -- (1281, 5) -- (8128, 0) -- (2809, 1) 
                                  -- (2882, 2) -- (8281, 1) -- (8128, 3) -- (2839, 4) 
                                  -- (2850, 3) -- (5041, 1) -- (4187, 2) 
                                                            -- (4141, 4) -- (4187, 2) 
                                               -- (5017, 2) -- (1764, 1) -- (6426, 4) 
                                                            -- (1782, 4) -- (8281, 1) 
                                  -- (2839, 4) -- (3969, 1) -- (6902, 2) 
                                                            -- (6903, 3) 
                     -- (8100, 1) 
                     -- (8177, 2) -- (7750, 0) -- (5041, 1) -- (4186, 3) -- (8614, 4) 
                                                            -- (4141, 4) -- (4186, 3) 
                                  -- (7744, 1) -- (4465, 0) 
                                  -- (7756, 4) -- (5671, 0) -- (7140, 3) -- (4096, 1) 
                                               -- (5625, 1) -- (2556, 0) 
                                                            -- (2556, 3) -- (5671, 0) 
                     -- (8128, 3) -- (2850, 0) -- (5041, 1) -- (4187, 2) 
                                                            -- (4141, 4) -- (4187, 2) 
                                               -- (5017, 2) -- (1764, 1) -- (6426, 4) 
                                                            -- (1782, 4) -- (8281, 1) 
                                  -- (2809, 1) 
                                  -- (2882, 2) -- (8256, 0) -- (5625, 1) -- (2512, 4)
                                                            -- (5688, 4) 
                                               -- (8281, 1) 
                                  -- (2839, 4)