Murder puzzle using or tools

murder.py

from ortools.sat.python import cp_model

persons = ['robert', 'john', 'george', 'yolando', 'christine', 'barbara']
rooms   = ['kitchen', 'bathroom', 'dining', 'living', 'pantry', 'study']
weapons = ['bag', 'firearm', 'gas', 'knife', 'poison', 'rope']

robert, john, george, yolando, christine, barbara = 0, 1, 2, 3, 4, 5
kitchen, bathroom, dining, living, pantry, study = 0, 1, 2, 3, 4, 5
bag, firearm, gas, knife, poison, rope = 0, 1, 2, 3, 4, 5

def gen6x6():
    return [[0 for _ in range(6)] for _ in range(6)]

A = gen6x6() # A[person][room]   is whether the person is in the room
B = gen6x6() # B[weapon][person] is whether the weapon is with the person
C = gen6x6() # C[weapon][room]   is whether the weapon is in the room

model = cp_model.CpModel()

# Define A, B, and C entries as decision variables

for i in range(6):
    for j in range(6):
        A[i][j] = model.NewIntVar(0, 1, f'{persons[i]}_in_{rooms[j]}')
        B[i][j] = model.NewIntVar(0, 1, f'{persons[j]}_has_{weapons[i]}')
        C[i][j] = model.NewIntVar(0, 1, f'{weapons[i]}_in_{rooms[j]}')

# All different constrains

# By rows
for i in range(6):
    x1, x2, x3 = None, None, None
    for j in range(6):
        if x1 is None:
            x1 = A[i][j]
            x2 = B[i][j]
            x3 = C[i][j]
        else:
            x1 = x1 + A[i][j]
            x2 = x2 + B[i][j]
            x3 = x3 + C[i][j]
    model.Add(x1 == 1)
    model.Add(x2 == 1)
    model.Add(x3 == 1)

# By cols
for j in range(6):
    x1, x2, x3 = None, None, None
    for i in range(6):
        if x1 is None:
            x1 = A[i][j]
            x2 = B[i][j]
            x3 = C[i][j]
        else:
            x1 = x1 + A[i][j]
            x2 = x2 + B[i][j]
            x3 = x3 + C[i][j]
    model.Add(x1 == 1)
    model.Add(x2 == 1)
    model.Add(x3 == 1)

# Consistency constrains

for r in range(6):
    for w in range(6):
        for p in range(6):
            # Any two entails the third so ABC can't be 011 or 101 or 110
            model.Add(A[p][r] * 4 + B[w][p] * 2 + C[w][r] != 6)
            model.Add(A[p][r] * 4 + B[w][p] * 2 + C[w][r] != 3)
            model.Add(A[p][r] * 4 + B[w][p] * 2 + C[w][r] != 5)

# first clue

"""
clue:
  1. The kitchen didn't have the rope, knife, bag, or firearm in it.
  2. The person in the kitchen is a man.
logic:
  !C[rope][kitchen]
  !C[knife][kitchen]
  !C[bag][kitchen]
  !C[firearm][kitchen]
  !A[barbara][kitchen]
  !A[christine][kitchen]
  !A[yolando][kitchen]
arithmetic:
  set to ones and zeros
"""

model.Add(C[rope][kitchen] == 0)
model.Add(C[knife][kitchen] == 0)
model.Add(C[bag][kitchen] == 0)
model.Add(C[firearm][kitchen] == 0)
model.Add(A[barbara][kitchen] == 0)
model.Add(A[christine][kitchen] == 0)
model.Add(A[yolando][kitchen] == 0)


# second clue

"""
clue:
  1. Yolando was either in the study or the bathroom and barbara was in the other.
  Note: To capture this using arithmetic ops, since every one is in a different room, we can
        assert that either yolando or barabar in the bathroom (x + y = 1) and the same for the study.
logic:
  A[barbara][bathroom] && A[yolando][study] || A[barbara][bathroom] && A[yolando][study]
arithmetic:
  A[barbara][bathroom] + A[yolando][bathroom] = 1
  A[barbara][study] + A[yolando][study] = 1
"""

model.Add(A[barbara][bathroom] + A[yolando][bathroom] == 1)
model.Add(A[barbara][study] + A[yolando][study] == 1)

# third clue

"""
clue:
  1. George didn't have the bag.
  2. Barbara didn't have the bag.
  3. The bag wasn't in the bathroom.
  4. The bag wasn't in the dining room.
logic:
  !B[bag][george]
  !B[bag][barbara]
  !C[bag][bathroom]
  !C[bag][dining]
arithmetic:
  same as logic just put zeros there
"""

model.Add(B[bag][george] == 0)
model.Add(B[bag][barbara] == 0)
model.Add(C[bag][bathroom] == 0)
model.Add(C[bag][dining] == 0)

# fourth clue

"""
clue:
  1. The rope was in the study room.
  3. The person in the study is not a man.
logic:
  !A[george][study]
  !A[robert][study]
  !A[john][study]
  C[rope][study]
arithmetic:
  same as logic just put zeros, or one
"""

model.Add(A[george][study] == 0)
model.Add(A[robert][study] == 0)
model.Add(A[john][study] == 0)
model.Add(C[rope][study] == 1)

# fifth clue

"""
clue:
  1. Either john or george was in the living room.
  Note: also using the difference constrain we can state the expression as in the arithmetic section.
logic:
  A[john][living] || A[george][living]
arithmetic:
  A[john][living] + A[george][living] >= 1
"""

model.Add(A[john][living] + A[george][living] >= 1)

# sixth clue

"""
clue:
  1. The knife wasn't in the dining room.
logic:
  !C[knife][dining]
arithmetic:
  zero
"""

model.Add(C[knife][dining] == 0)

# seventh clue

"""
clue:
  1. yolando wasn't in the study.
  2. yolando wan't in the pantry.
logic:
  !A[yolando][study]
  !A[yolando][pantry]
arithmetic:
  zero
"""

model.Add(A[yolando][study] == 0)
model.Add(A[yolando][pantry] == 0)

# eigth clue

"""
clue:
  1. The firearm was with george.
logic:
  B[firearm][george]
arithmetic:
  one
"""

model.Add(B[firearm][george] == 1)

# Gassed in the pantry clue

"""
clue:
  1. The victim was gassed.
logic:
  C[gas][pantry]
arithmetic:
  one
"""

model.Add(C[gas][pantry] == 1)

# Solution printing
class VarArraySolutionPrinter(cp_model.CpSolverSolutionCallback):
    """Print intermediate solutions."""

    def __init__(self):
        cp_model.CpSolverSolutionCallback.__init__(self)
        self.__solution_count = 0

    def on_solution_callback(self):
        self.__solution_count += 1
        for p in range(6):
            for r in range(6):
                if self.Value(A[p][r]):
                    print(f"{persons[p]} is in {rooms[r]}")
        for p in range(6):
            for w in range(6):
                if self.Value(B[w][p]):
                    print(f"{persons[p]} has the {weapons[w]}")
        for p in range(6):
            if self.Value(A[p][pantry]):
                print(f"The murderer is {persons[p]}")
        print()

    def solution_count(self):
        return self.__solution_count

solver = cp_model.CpSolver()
status = solver.Solve(model)
solution_printer = VarArraySolutionPrinter()
status = solver.SearchForAllSolutions(model, solution_printer)