sat solver example

magic_square.py

from ortools.sat.python import cp_model

def run():
    model = cp_model.CpModel()
    vars = []
    for i in range(9):
        vars.append(model.NewIntVar(1, 9, str(i)))

    model.AddAllDifferent(vars)

    for col in range(3):
        model.Add(vars[col]+vars[col+3]+vars[col+6]==15)

    for row in range(3):
        idx = row*3
        model.Add(vars[idx]+vars[idx+1]+vars[idx+2]==15)

    model.Add(vars[0]+vars[4]+vars[8]==15)
    model.Add(vars[2]+vars[4]+vars[6]==15)

    solver = cp_model.CpSolver()
    status = solver.Solve(model)

    if status==cp_model.FEASIBLE:
        for row in range(3):
            idx = row*3
            print("{0} {1} {2}".format(*[solver.Value(vars[i]) for i in range(idx, idx+3)]))
    else:
        print(f"Unexpected status: {status}")

run()