Antwoord op logica puzzle

antwoord.py

import time
start_time = time.time()
from z3 import *
opties = {
    'nummer': ['1', '2', '3', '4', '5'],
    'kleur': ['Rood', 'Groen', 'Geel', 'Blauw', 'Wit'],
    'dier': ['Honden', 'Vogels', 'Katten', 'Paarden', '[Vis]'],
    'drinkt': ['Thee', 'Koffie', 'Melk', 'Bier', 'Water'],
    'rookt': ['PallMall', 'Dunhill', 'Blends', 'BlueMasters', 'Prince'],
    'nationaliteit': ['Brit', 'Zweed', 'Deen', 'Noor', 'Duitser'],
}
parameters = {k:[Int(f"{k}_{i}") for i in range(len(v))] for k,v in opties.items()}
all_parameters = {}
for optie, params in parameters.items():
    for p in params:
        all_parameters[str(p)] = p
def p(optie, index):
    return all_parameters[f"{optie}_{index}"]
def o(optie, waarde):
    return opties[optie].index(waarde)

def p_is_o(optie, index, waarde):
    return p(optie, index) == o(optie, waarde)
def p_isnot_o(optie, index, waarde):
    return p(optie, index) != o(optie, waarde)
def if_than_extended(if_optie, if_index, if_waarde, true_condition):
    return If(p_is_o(if_optie, if_index, if_waarde), true_condition, p_isnot_o(if_optie, if_index, if_waarde))
def if_than(if_optie, if_index, if_waarde, than_optie, than_index, than_waarde):
    return if_than_extended(if_optie, if_index, if_waarde, p_is_o(than_optie, than_index, than_waarde))
def heeft_buurman(if_optie, if_index, if_waarde, buurman_optie, buurman_waarde):
    if 0 < if_index < 4:
        true_con = Or(p_is_o(buurman_optie, if_index-1, buurman_waarde), p_is_o(buurman_optie, i+1, buurman_waarde))
    elif i == 0:
        true_con = p_is_o(buurman_optie, i+1, buurman_waarde)
    elif i == 4:
        true_con = p_is_o(buurman_optie, i-1, buurman_waarde)
    return if_than_extended(if_optie, if_index, if_waarde, true_con)


s = Solver()
# Basis regels
# Huisnummers staan vast 
for i, pw in enumerate(parameters['nummer']):
    s.add(pw == i) 
for p_naam, p_waardes in parameters.items():
    # Per categorie zijn de waardes unique
    s.add(Distinct(p_waardes))
    # Alle waardes liggen tussen 0 en 5
    for pw in p_waardes:
        s.add(And(pw >= 0, pw < 5))

# Regels
# Brit woont in rood huis
for i in range(5):
    s.add(if_than('nationaliteit', i, 'Brit', 'kleur', i, 'Rood'))
# Zweed heeft honden
for i in range(5):
    s.add(if_than('nationaliteit', i, 'Zweed', 'dier', i, 'Honden'))
# Deen drinkt thee
for i in range(5):
    s.add(if_than('nationaliteit', i, 'Deen', 'drinkt', i, 'Thee'))
# Het groene huis staat direct links from het witte huis
for i in range(1,5):
    s.add(if_than('kleur', i, 'Wit', 'kleur', i-1, 'Groen'))
# Het bewoner van het groen huis drinkt koffie
for i in range(5):
    s.add(if_than('kleur', i, 'Groen', 'drinkt', i, 'Koffie'))
# Het bewoner van het groen huis drinkt koffie
for i in range(5):
    s.add(if_than('rookt', i, 'PallMall', 'dier', i, 'Vogels'))
# Het bewoner van het gele huis rookt Dunhull
for i in range(5):
    s.add(if_than('kleur', i, 'Geel', 'rookt', i, 'Dunhill'))
# Het bewoner van het middelste huis drinkt melk
s.add(p_is_o('drinkt', 2, 'Melk'))
# De Noor woont in het eerste huis
s.add(p_is_o('nationaliteit', 0, 'Noor'))
# De bewoner die Blends rookt woont naast de bweoner die katten heeft
for i in range(5):
    s.add(heeft_buurman('rookt', i, 'Blends', 'dier', 'Katten'))
# De bewoner die paarden heeft woont naast de bewoner die Dunhill rookt
for i in range(5):
    s.add(heeft_buurman('dier', i, 'Paarden', 'rookt', 'Dunhill'))
# De bewoner die Bluemasters rookt drinkt Bier
for i in range(5):
    s.add(if_than('rookt', i, 'BlueMasters', 'drinkt', i, 'Bier'))
# De duister rookt Prince
for i in range(5):
    s.add(if_than('nationaliteit', i, 'Duitser', 'rookt', i, 'Prince'))
# De Noor woont naast het blauwe huis
for i in range(5):
    s.add(heeft_buurman('nationaliteit', i, 'Noor', 'kleur', 'Blauw'))
# De bewoner die Blends rookt woont naast de bewoner die water drinkt
for i in range(5):
    s.add(heeft_buurman('rookt', i, 'Blends', 'drinkt', 'Water'))

# Oplossing
s.check()
solution = s.model()

# Print uitkomst
solution_txt = [opties.copy() for i in range(5)]
for p_naam, p_waardes in parameters.items():
    for pw in p_waardes:
        wat, huisnummer = str(pw).split('_')
        keuze = opties[wat][int(str(solution[pw]))]
        solution_txt[int(huisnummer)][wat] = keuze
rows = []
for optie in opties:
    rows.append([optie] + [str(solution_txt[i][optie]) for i in range(5)])
row_format ="{:>15}" * (len(rows))
for row in rows:
    print(row_format.format(*row))
end_time = time.time()
print(f'Berekening duurde {end_time - start_time} seconden.')