Sudoku Solver using MILP

sudoku_solver_milp.py

from ortools.linear_solver import pywraplp
import numpy as np

def solve(sudoku):
    """
    Solve the sudoku passed
    :param sudoku: a 9x9 sudoku field
    :return: the completed 9x9 sudoku field
    """
    # Create the model.
    solver_name = 'SCIP'
    model = pywraplp.Solver.CreateSolver(solver_name)
    model.EnableOutput()

    # create clues parameter
    c = convert_sudoku_to_binary_array(sudoku)

    # create index sets
    I, J, K, P, Q = list(range(9)), list(range(9)), list(range(9)), list(range(3)), list(range(3))

    # create decision vars
    X = {}
    for i in I:
        for j in J:
            for k in K:
                X[(i, j, k)] = model.BoolVar('task_i%ij%ik%i' % (i, j, k))

    # constraint 0: Each cell must contain a digit
    for i in I:
        for j in J:
            model.Add(sum(X[(i, j, k)] for k in K) == 1)

    # constraint 1: Each row must contain each digit exactly once
    for i in I:
        for k in K:
            model.Add(sum(X[(i, j, k)] for j in J) == 1)

    # constraint 2: Each column must contain each digit exactly once
    for j in J:
        for k in K:
            model.Add(sum(X[(i, j, k)] for i in I) == 1)

    # constraint 3: Each subgrid must contain each digit exactly once
    for p in P:
        # calculate i_subset from p for sum-formula
        i_subset = list(range(3 * p, 3 * p + 3))
        for q in Q:
            # calculate j_subset from q for sum-formula
            j_subset = list(range(3 * q, 3 * q + 3))
            for k in K:
                model.Add(sum(X[(i, j, k)] for i in i_subset for j in j_subset) == 1)

    # constraint 4: Comply with the clues provided
    for i in I:
        for j in J:
            for k in K:
                model.Add(X[(i, j, k)] >= c[(i, j, k)])

    # Zielfunktion
    model.Minimize(0)

    status = model.Solve()
    if status == pywraplp.Solver.OPTIMAL:
        model_solved = True
        solution = convert_x_to_sudoku(X)
    else:
        model_solved = False
        solution = sudoku

    return model_solved, solution


def convert_sudoku_to_binary_array(sudoku):
    """
    Takes a 2x2 sudoku np-array and derives a 3-d clue param by adding a dimension
    of binary flags for each digit as required by the model

    :param sudoku: the sudoku np-array
    :return: the clue parameter for the model
    """

    binary_sudoku = np.zeros((9, 9, 9), dtype=int)
    for i in range(9):
        for j in range(9):
            if sudoku[i, j] != 0:
                binary_sudoku[i, j, sudoku[i, j] - 1] = 1

    return binary_sudoku


def convert_x_to_sudoku(X):
    """
    Convert the solver solution back to a sudoku 2x2 field

    :param X: the solution variable x[i,j,k]
    :return: a sudoku 2x2 field
    """
    sudoku = np.zeros(shape=(9, 9), dtype=np.int8)
    for i in range(0, 9):
        for j in range(0, 9):
            for k in range(0, 9):
                if X[(i, j, k)].solution_value() == 1:
                    sudoku[i, j] = k + 1
    return sudoku


if __name__ == '__main__':
    # Define a sudoku; 0 = undefined, 1-9 are clues
    sudoku = np.array([[0, 7, 2, 4, 8, 5, 0, 0, 0],
                       [4, 0, 8, 2, 0, 0, 0, 0, 0],
                       [5, 0, 0, 0, 0, 9, 4, 0, 0],
                       [0, 0, 5, 0, 0, 1, 0, 0, 8],
                       [0, 0, 0, 0, 6, 0, 0, 0, 0],
                       [1, 0, 0, 5, 0, 0, 9, 0, 0],
                       [0, 0, 4, 1, 0, 0, 0, 0, 5],
                       [0, 0, 0, 0, 0, 4, 3, 0, 7],
                       [0, 0, 0, 7, 3, 8, 2, 1, 0]])

    solved, solution = solve(sudoku)
    if solved:
        print(solution)
    else:
        print("No solution found. Unfeasible instance provided.")