Magic Square (Mine) - 魔方陣

どんな問題？

Magic Square
http://www.checkio.org/mission/magic-square/

魔方陣を完成させよ。

魔方陣とは、１からマスの数までの数字を、数字が重複せず、かつ縦横斜め全ての並びの合計値が同じになるように配置したもの。

引数は、未完成の魔方陣が list of list で渡される。空のマスは0となる。
戻り値は、引数の0に数字を入れて返す。

魔方陣のサイズは、3x3、4x4、5x5のいずれか。

例題：

checkio([
    [2, 7, 6],
    [9, 5, 1],
    [4, 3, 0]
]) == [[2, 7, 6], [9, 5, 1], [4, 3, 8]]
 
checkio([
    [0, 0, 0],
    [0, 5, 0],
    [0, 0, 0]
]) ==  [[2, 7, 6], [9, 5, 1], [4, 3, 8]]
#  or  [[4, 9, 2], [3, 5, 7], [8, 1, 6]]
 
checkio([[1, 15, 14, 4],
         [12, 0, 0, 9],
         [8, 0, 0, 5],
         [13, 3, 2, 16]
]) == [[1, 15, 14, 4], [12, 6, 7, 9], [8, 10, 11, 5], [13, 3, 2, 16]]

どうやって解く？

いつものようにグラフ探索(DFS)で解けそうだ。ボードの状態を頂点として、空きマスに使われていない数字を入れていく。

まず、Wikipediaによると、魔方陣の一列の和は、

となるらしい。
最初に、一列の和を求めておく。

size = len(data)  # 辺の長さ
line_sum = size * (size ** 2 + 1) / 2  # 一列の和

次に、ボードが矛盾していないか判定する関数を書いた。
縦横斜めの数字が一致するかどうかをチェックする。0が含まれる列は、チェックから除外している。

def is_valid(b):
    lines = []
    for i in range(size):
        # 横
        lines.append([b[idx] for idx in [i * size + x for x in range(size)]])
        # 縦
        lines.append([b[idx] for idx in [y * size + i for y in range(size)]])
    # 斜め
    lines.append([b[idx] for idx in [i * (size + 1) for i in range(size)]])
    lines.append([b[idx] for idx in [(i + 1) * (size - 1) for i in range(size)]])
    # 全ての列の合計値をチェック
    return all(0 in line or sum(line) == line_sum for line in lines)

ボードの状態と未使用の数字をタプルにしてスタックに追加する。
未使用の数字を空きマスに入れて新しいボードを作成し、ボードに矛盾がなければ、スタックに追加する。
これを未使用の数字が無くなるまで繰り返す。
（再帰を使っても良かったのだが、いつものクセでstackを使った。）

まとめ

def checkio(data):
    def is_valid(b):
        lines = []
        for i in range(size):
            # horizontal
            lines.append([b[idx] for idx in [i * size + x for x in range(size)]])
            # vertical
            lines.append([b[idx] for idx in [y * size + i for y in range(size)]])
        # diagonal
        lines.append([b[idx] for idx in [i * (size + 1) for i in range(size)]])
        lines.append([b[idx] for idx in [(i + 1) * (size - 1) for i in range(size)]])
        return all(0 in line or sum(line) == line_sum for line in lines)

    size = len(data)
    line_sum = size * (size ** 2 + 1) / 2

    board = [n for row in data for n in row]
    unused = {i + 1 for i in range(size ** 2)} - set(board)  # 未使用の数字
    stack = [(board, unused)]
    while stack:
        board, unused = stack.pop()
        if not unused:  # 未使用の数字が無くなれば終了
            return [board[i * size:i * size + size] for i in range(size)]
        for num in unused:  # 未使用の数字を空きマス(0)に入れて新しいボードを作成
            new_board = board[:]
            new_board[board.index(0)] = num
            if is_valid(new_board):  # ボードが有効ならスタックに追加
                stack.append((new_board, unused - {num}))

http://www.checkio.org/mission/magic-square/publications/natsuki/python-3/first/