数独を解くNuSMVソースコードのジェネレータ

gensmv.py

#!/usr/bin/env python3

import sys

#-----------------------------------------------------#

def get_row_set(sudoku, r_idx, c_idx):
    return set([i for i in sudoku[r_idx] if sudoku[r_idx][c_idx] != i])

def get_col_set(sudoku, r_idx, c_idx):
    return set([sudoku[i][c_idx] for i in range(9) if r_idx != i])

def get_sq_set(sudoku, r_idx, c_idx):
    base_row = (r_idx // 3) * 3
    base_col = (c_idx // 3) * 3
    res = []
    for i in range(3):
        cur_row = base_row + i
        for j in range(3):
            cur_col = base_col + j
            res.append(sudoku[cur_row][cur_col])
    return set(res)

def get_possible_list(sudoku, r_idx, c_idx):
    row_set = get_row_set(sudoku, r_idx, c_idx)
    col_set = get_col_set(sudoku, r_idx, c_idx)
    sq_set = get_sq_set(sudoku, r_idx, c_idx)
    forbid_set = row_set | col_set | sq_set
    return [j for j in range(10) if not (j in forbid_set)]

def get_possible_list_str(sudoku, r_idx, c_idx):
    possible_list = get_possible_list(sudoku, r_idx, c_idx)
    return [str(i) for i in possible_list]

def gen_addeq_row(row_char, col_idx):
    keys = []
    for i in range(9):
        if col_idx == i:
            continue
        keys.append("{}{}".format(row_char, i))
    return " + ".join(keys)

def gen_addeq_col(row_char, col_idx):
    keys = []
    row_idx = ord(row_char) - ord("A")
    for i in range(9):
        if row_idx == i:
            continue
        keys.append("{}{}".format(chr(i + ord("A")), col_idx))
    return " + ".join(keys)

def gen_addeq_sq(row, idx):
    base_row = ord("A") + ((ord(row)-ord("A")) // 3) * 3
    base_col = 0 + (idx // 3) * 3
    res = []
    for i in range(3):
        cur_row = chr(base_row + i)
        for j in range(3):
            cur_col = base_col + j
            if cur_row == row and cur_col == idx:
                continue
            res.append("{}{}".format(cur_row, cur_col))
    return " + ".join(res)


def spec_row_eq():
    eqs = []
    for i in range(9):
        row = chr(ord("A") + i)
        eq = ["{}{}".format(row, j) for j in range(9)]
        eqs.append("{} = 45".format(" + ".join(eq)))
    return " &\n".join(eqs)

def spec_col_eq():
    eqs = []
    for i in range(9):
        eq = ["{}{}".format(chr(ord("A")+j), i) for j in range(9)]
        eqs.append("{} = 45".format(" + ".join(eq)))
    return " &\n".join(eqs)

def spec_sq_eq():
    eqs = []
    for i in range(3):
        base_row = ord("A") + i * 3
        for j in range(3):
            base_col = j * 3
            eq = []
            for k in range(3):
                row = chr(base_row + k)
                for x in range(3):
                    col = base_col + x
                    eq.append("{}{}".format(row, col))
            eqs.append("{} = 45".format(" + ".join(eq)))
    return " &\n".join(eqs)

def spec_row_neq(sudoku, blanks):
    neqs = []
    for r, c in blanks:
        row_char = chr(ord("A") + r)
        key = "{}{}".format(row_char, c)
        for i in range(c+1, 9):
            if sudoku[r][i] != 0:
                continue
            neqkey = "{}{}".format(row_char, i)
            neqs.append("!({} = {})".format(key, neqkey))
    return " &\n".join(neqs)

def spec_col_neq(sudoku, blanks):
    neqs = []
    for r, c in blanks:
        row_char = chr(ord("A") + r)
        key = "{}{}".format(row_char, c)
        for i in range(r+1, 9):
            if sudoku[i][c] != 0:
                continue
            neqkey = "{}{}".format(chr(ord("A")+i), c)
            neqs.append("!({} = {})".format(key, neqkey))
    return " &\n".join(neqs)


def spec_sq_neq():
    eqs = []
    for i in range(3):
        base_row = ord("A") + i * 3
        for j in range(3):
            base_col = j * 3
            eq = []
            keys = []
            for k in range(3):
                row = chr(base_row + k)
                for x in range(3):
                    col = base_col + x
                    keys.append("{}{}".format(row, col))
            for k in range(len(keys)):
                for x in range(k+1, len(keys)):
                    eqs.append("!({} = {})".format(keys[k], keys[x]))
    return " &\n".join(eqs)

#-----------------------------------------------------#
def fill_easy_blank(sudoku):
    for r in range(9):
        for c in range(9):
            v = sudoku[r][c]
            if v != 0:
                continue
            p = get_possible_list(sudoku, r, c)
            if len(p) == 1:
                sudoku[r][c] = p[0]

def sum_sudoku(sudoku):
    s = 0
    for row in sudoku:
        s += sum(row)
    return s

def print_sudoku(sudoku):
    for row in sudoku:
        print(" ".join([str(i) for i in row]))

#-----------------------------------------------------#

def print_define_const(sudoku):
    print("DEFINE")
    for i, row in enumerate(sudoku):
        r = chr(ord("A") + i)
        row_set = set(row)
        for c, v in enumerate(row):
            key = "{}{}".format(r, c)
            if v != 0:
                print("  {} := {};".format(key, v))

def print_var(sudoku):
    print("VAR")
    for i, row in enumerate(sudoku):
        r = chr(ord("A") + i)
        row_set = set(row)
        for c, v in enumerate(row):
            key = "{}{}".format(r, c)
            if v != 0:
                continue
            else:
                possible_list = get_possible_list_str(sudoku, i, c)
                possible_list.append("0")
                print("  {} : {{{}}};".format(key, ", ".join(sorted(possible_list))))

def print_assign(sudoku):
    print("ASSIGN")
    for i, row in enumerate(sudoku):
        r = chr(ord("A") + i)
        row_set = set(row)
        for c, v in enumerate(row):
            key = "{}{}".format(r, c)
            if v == 0:
                print("  init({}) := 0;".format(key))
                print("  next({}) := case".format(key))
                #addeq_row = gen_addeq_row(r, c)
                #addeq_col = gen_addeq_col(r, c)
                #addeq_sq = gen_addeq_sq(r, c)
                possible = set(get_possible_list(sudoku, i, c))
                #for j in range(1, 10):
                #    if j in possible:
                #        print("    {} = {} | {} = {} | {} = {} : {};".format(addeq_row, 45-j, addeq_col, 45-j, addeq_sq, 45-j, j))
                possible = get_possible_list_str(sudoku, i, c)
                print("    TRUE : {{{}}};".format(", ".join(possible)))
                print("  esac;\n")

def print_spec(sudoku, blank_key):
    print("SPEC")
    print("!EF(")
    tmp = []
    #row_eq = spec_row_eq()
    #col_eq = spec_col_eq()
    #sq_eq = spec_sq_eq()
    tmp.append(spec_row_neq(sudoku, blank_key))
    tmp.append(spec_col_neq(sudoku, blank_key))
    tmp.append(spec_sq_neq())
    print(" &\n".join(tmp))

    #col_eq = spec_col_neq()
    #sq_neq = spec_sq_neq()
    print(")")

#-----------------------------------------------------#

def usage():
    print("usage: {} <sudoku.csv>".format(sys.argv[0]))
    sys.exit(1)

def main():
    if len(sys.argv) != 2:
        usage()
    sudoku = []
    blank_key = []
    with open(sys.argv[1]) as f:
        r = 0
        for line in f:
            line = line.strip().split(" ")
            row = []
            for c, v in enumerate(line):
                v = int(v)
                row.append(v)
                if v == 0:
                    blank_key.append((r, c))
            sudoku.append(row)
            r += 1

    #prev_sum = sum_sudoku(sudoku)
    #while True:
    #    fill_easy_blank(sudoku)
    #    cur_sum = sum_sudoku(sudoku)
    #    if cur_sum == prev_sum:
    #        break
    #    prev_sum = cur_sum

    #if sum_sudoku(sudoku) == 45 * 9:
    #    print("*** solved easily ***")
    #    print_sudoku(sudoku)
    #    return

    print("MODULE main\n")
    print_define_const(sudoku)
    print("")
    print_var(sudoku)
    print("")
    print_assign(sudoku)
    print("")
    print_spec(sudoku, blank_key)

#-----------------------------------------------------#
if __name__ == "__main__":
    main()