Rubyで魔方陣

知人に6x6の魔方陣を見せられて、そういえば魔方陣ってやったことないなと思ってRubyで解いてみた。モヤモヤと悩みつつ、良い練習になった。

RubyにはEnumerable#permutationという強力なのがある。多重ループを使わなくても、「(0..3).to_a.permutation」とかやれば組み合わせの集合を一気に生成できる。でも、せっかくなら素数の集合をいろいろ突っ込むと答えが出るようなやつで、3x3だけじゃなくて、任意の次数に対応する汎用的なソルバーを作りたいし、習作としてはRubyっぽくやってみたいと思って、以下のように書いた。左上から順に数字を入れて、DFSもしくはBFS。

require './core_ext.rb'
require './matrix_utils.rb'

class Matrix
  attr_accessor :items, :matrix, :order
  include MatrixUtils

  def initialize(arg)
    case arg
    when Array then arr = arg
    when Range then arr = arg.to_a
    end

    raise "Not a squre" unless arr.size.square_number?

    @order = Math.sqrt(arr.size).to_i
    @items  = arr
    @matrix = Array.new(@order**2, 0)
  end

  def self.duplicate(arg)
    Marshal.load(Marshal.dump(arg))
  end

  def next_blank_position
    p = matrix.take_while{|i| i != 0}.size
    p != matrix.size ? p : false
  end

  def candidate_numbers(cell)
    left = items - matrix
    x, y = coordinate(cell)
    case
    when x == @order -1
      num = magic_constant - sum_cells(row_indices(y))
      left.include?(num) ? [num] : nil
    when y == @order -1
      num = magic_constant - sum_cells(column_indices(x))
      left.include?(num) ? [num] : nil
    else
      left.select { |i|
        sum_cells(row_indices(y)) + i    <= magic_constant &&
        sum_cells(column_indices(x)) + i <= magic_constant
      }
    end
  end

  def symmetry_constraint_fullfilled?
    if top_left * top_right * buttom_left * buttom_right == 0
      true
    else
      top_left > top_right &&
      top_left > buttom_left &&
      top_left > buttom_right &&
      top_right > buttom_left
    end
  end

  def magic?
    return false unless sum_cells(diagonal_indices1) == magic_constant
    return false unless sum_cells(diagonal_indices2) == magic_constant

    0.upto(@order - 1) do |i|
      if sum_cells(column_indices(i)) != magic_constant ||
        sum_cells(column_indices(i)) != magic_constant
        return false
      end
    end
    true
  end
end

class Solver
  def initialize(matrix)
    @matrix = matrix
  end

  def solve
    if @matrix.next_blank_position
      stack = place_a_new_number_in_the_next_blank_cell(@matrix)
      while !stack.empty?
        matrix = Matrix.duplicate(stack.shift)
        Solver.new(matrix).solve
      end
    else
      if @matrix.magic?
        p @matrix if @matrix.magic?
      else
      end
    end
  end

  def place_a_new_number_in_the_next_blank_cell(matrix)
    stack = []
    if cell = matrix.next_blank_position
      nums = matrix.candidate_numbers(cell)
      return stack unless nums
      nums.each do |e|
        new_matrix = Matrix.duplicate(matrix)
        new_matrix.matrix[cell] = e
        stack << new_matrix if new_matrix.symmetry_constraint_fullfilled?
      end
    end
    stack
  end
end

if  __FILE__ == $0
#  m = Matrix.new([17, 113, 47, 89,  59,  29, 71, 5, 101])
#  m = Matrix.new([8, 1, 6, 3, 5, 7, 4, 9, 2])
#  m = Matrix.new((1..16))
#  m[0,0] = 16
  m = Matrix.new((1..9))
  Solver.new(m).solve
  exit
end

途中でMatrixクラスが大きくなったので、Solverから直接見えなくてもいいような抽象度の低いところをmoduleで分けてみた。

module MatrixUtils
  def magic_constant
    @magic_constant ||= @items.inject(:+) / @order
  end

  def [](x, y)
    @matrix[x + y * @order]
  end

  def []=(x, y, n)
    @matrix[x + y * @order] = n
  end

  def top_left; @matrix[0]; end
  def top_right; @matrix[@order - 1]; end
  def buttom_left; @matrix[@order * (@order - 1)]; end
  def buttom_right; @matrix[@order**2 - 1]; end

  def column_indices(x)
    (0..(@order - 1)).map{|i| i * @order + x}
  end

  def row_indices(y)
    (0..(@order - 1)).map{|i| i + @order * y}
  end

  def diagonal_indices1
    (0..(@order - 1)).map{|i| (@order + 1) * i}
  end

  def diagonal_indices2
    (1..@order).map{|i| (@order - 1) * i}
  end

  def sum_cells(indices)
    indices.inject(0) {|sum, i| sum += @matrix[i]}
  end

  def coordinate(cell)
    x = cell % @order
    y = cell / @order
    return x, y
  end

  def col_width
    @col_width ||= @items.max.to_s.size
  end

  def to_s
    @matrix.map! { |e| e ? e : 0}
    @matrix.each_slice(@order) do |arr|
      print "|"
      arr.each do |e|
        print "%#{col_width}d|" % e
      end
      puts
    end
  end
end

この辺はRSpecも書いた。

require './666.rb'

describe Numeric do
  it "should return true if n is a square number" do
    [4, 9, 16, 25, 36].all? {|n| n.square_number? }.should == true
  end

  it "should return false if n is not a square number" do
    [5, 9, 16, 25, 36].all? {|n| n.square_number? }.should == false
  end
end

describe Matrix do
  before(:each) do
    @m = Matrix.new((1..9))
    @m.matrix = @m.items.dup
  end

  describe "initialize" do
    context "with Range object" do
      it "should be initialized" do
        @m = Matrix.new((1..9))
        @m.matrix = @m.items.dup
        @m.items.should have(9).items
        @m.matrix.should have(9).items
      end
    end

    context "with Array object" do
      it "should be initialized" do
        @m = Matrix.new((1..9).to_a)
        @m.matrix = @m.items.dup
        @m.items.should have(9).items
        @m.matrix.should have(9).items
      end
    end

    context "with Matrix object" do
      it "should be initialized" do
        @m2 = Matrix.duplicate(@m)
        @m2.should be_an_instance_of Matrix
        @m2.matrix.should == @m2.matrix
      end
    end
  end

  it "should have a magic_constant of 15" do
    @m.magic_constant.should == 15
  end

  it "should return 1 on the top-right cell" do
    @m[0, 0].should == 1
  end

  it "should return 5 for the center cell" do
    @m[1, 1].should == 5
  end

  it "should sum up column cells" do
    @m.sum_cells(@m.column_indices(0)).should == 12
    @m.sum_cells(@m.column_indices(1)).should == 15
    @m.sum_cells(@m.column_indices(2)).should == 18
  end

  it "should sum up row cells" do
    @m.sum_cells(@m.row_indices(0)).should == 6
    @m.sum_cells(@m.row_indices(1)).should == 15
    @m.sum_cells(@m.row_indices(2)).should == 24
  end

  it "should not be a magic matrix" do
    @m.should_not be_magic
  end

  describe "3x3 magic matrix" do
    before do
      @m = Matrix.new([8, 1, 6, 3, 5, 7, 4, 9, 2])
      @m.matrix = @m.items
      p @m
    end

    it "should be a magic matrix" do
      @m.magic?.should be_true
      @m.should be_magic
    end
  end
end

平方数を判定するメソッドがなかったので、以下のように標準クラスに定義。

class Numeric
  def square_number?
    Math.sqrt(self).integer?
  end
end

class Float
  def integer?
    self.ceil == self
  end
end

実行結果は、

$ time ruby 666.rb
|8|1|6|
|3|5|7|
|4|9|2|

ruby 666.rb  0.12s user 0.01s system 31% cpu 0.396 total

という感じ。1から9を使う魔方陣は対照なものを除くと、この1つだけ。あらかじめ魔方陣が作れることがわかっている素数の9個の組み合わせを入れると、

|101|  5| 71|
| 29| 59| 89|
| 47|113| 17|

と答えが出る。

Rubyではifの条件式の意図が分かりづらかったら、意図が分かる名前のメソッドを作って、そっちの式を移せというような文化があるように思う。これを突き詰めると、コメントを入れなくてもセルフドキュメンテッドなソースコードになっていくような気もする。例えば、

      case
      when x == @order -1
        num = magic_constant - sum_cells(row_indices(y))
        left.include?(num) ? [num] : nil
      when y == @order -1
        num = magic_constant - sum_cells(column_indices(x))
        left.include?(num) ? [num] : nil
      else
        left.select { |i|
          sum_cells(row_indices(y)) + i    <= magic_constant &&
          sum_cells(column_indices(x)) + i <= magic_constant
        }
      end

というのがある。初めて見る人や、1年後に自分自身でこのコードを読むことを考えたら、

    def edge
      @edge ||= @order - 1
    end
    :
    :
      case
      when x == edge
        num = magic_constant - sum_cells(row_indices(y))
        left.include?(num) ? [num] : nil
      when y == edge
        num = magic_constant - sum_cells(column_indices(x))
        left.include?(num) ? [num] : nil
      else
        left.select { |i|
          sum_cells(row_indices(y)) + i    <= magic_constant &&
          sum_cells(column_indices(x)) + i <= magic_constant
        }
      end

と、@order - 1に名前を付けたほうがいいのかもしれない。うーん。

せっかくn次の魔方陣が調べられるようにしたはいいけど、4次ですら答えがチョロ……、チョロ……としか返ってこない。左上の数字1つを固定してみたら20分ほどでチェックが終了。あれこれ解が出てきた。最初に気付けよという感じだけど、計算量はO((n^2)!)と、組み合わせが爆発するのだった。16!は10進で14桁にもなる。25!だと26桁。

最初、全ての数字の組み合わせをセルに突っ込むように書いたら、3次の魔方陣を調べ終わるのに30秒ほどかかった。これはO((n^2)^2)で、さすがにひどいので、それぞれの数字は1回しか使えないという条件を付けると、2、3秒に縮んだ。これでO((n^2)!)

さらに、セルを埋める途中でも縦横の数字合計を随時チェックする枝刈り（candidate_numbes）を入れたら、これで0.2秒に縮んだ。回転対称と鏡像対象を省くようにして、最終的に0.12秒。しかし、それでも計算量のオーダーは変わらず、5次や6次はとても無理だ。

Wikipediaを見ると魔方陣の構成法というのは書いてあるけど、一般的な発見方法というのは書いていない。

遺伝的アルゴリズムで発見する方法というのに言及があって、それは面白そう。Matrixクラスが手元にあるので、試すのも簡単そう。