test.py

from z3 import *
import json
import sys

# Define all categories in a single list of tuples:
# (House is implicit; each row's first column will be the house number.)
categories = [
    ("Name", ["Bob", "Arnold", "Carol", "Alice", "Peter", "Eric"]),
    ("Mother", ["Sarah", "Janelle", "Aniya", "Kailyn", "Holly", "Penny"]),
    ("Children", ["Fred", "Samantha", "Bella", "Meredith", "Alice", "Timothy"]),
    ("Vacation", ["city", "mountain", "camping", "beach", "cruise", "cultural"]),
    ("BookGenre", ["romance", "mystery", "historical fiction", "science fiction", "biography", "fantasy"])
]

# No need to change here, automatically processing
NUM_POSITIONS = len(categories[0][1])

item_to_cat_and_index = {}
for cat_idx, (cat_str, item_list) in enumerate(categories):
    for item_idx, item_str in enumerate(item_list):
        item_to_cat_and_index[(cat_str, item_str)] = (cat_idx, item_idx)

Vars = []
for cat_idx, (cat_name, item_list) in enumerate(categories):
    var = IntVector(cat_name, len(item_list))
    Vars.append(var)

s = Solver()

for cat_idx, (cat_name, item_list) in enumerate(categories):
    for item_idx, item_str in enumerate(item_list):
        s.add(Vars[cat_idx][item_idx] >= 1, Vars[cat_idx][item_idx] <= NUM_POSITIONS)
    s.add(Distinct(Vars[cat_idx]))

def pos(cat_str, item_str):
    (cat_idx, item_idx) = item_to_cat_and_index[(cat_str, item_str)]
    return Vars[cat_idx][item_idx]

# All clues here
# 1. The person who loves beach vacations is not in the second house.
s.add(pos("Vacation", "beach") != 2)

# 2. The person who loves fantasy books is somewhere to the left of Peter.
s.add(pos("BookGenre", "fantasy") < pos("Name", "Peter"))

# 3. The person whose mother's name is Sarah is the person who prefers city breaks.
s.add(pos("Mother", "Sarah") == pos("Vacation", "city"))

# 4. The person who enjoys camping trips is somewhere to the right of Peter.
s.add(pos("Vacation", "camping") > pos("Name", "Peter"))

# 5. The person who likes going on cruises is the person's child is named Meredith.
s.add(pos("Vacation", "cruise") == pos("Children", "Meredith"))

# 6. There is one house between the person who is the mother of Timothy and Eric.
s.add(Abs(pos("Mother", "Timothy") - pos("Name", "Eric")) == 2)

# 7. The person whose mother's name is Janelle is not in the second house.
s.add(pos("Mother", "Janelle") != 2)

# 8. The person's child is named Fred is somewhere to the left of Eric.
s.add(pos("Children", "Fred") < pos("Name", "Eric"))

# 9. The person who goes on cultural tours is in the fourth house.
s.add(pos("Vacation", "cultural") == 4)

# 10. The person whose mother's name is Janelle is not in the first house.
s.add(pos("Mother", "Janelle") != 1)

# 11. The person whose mother's name is Holly is somewhere to the right of the person who loves historical fiction books.
s.add(pos("Mother", "Holly") > pos("BookGenre", "historical fiction"))

# 12. The person's child is named Bella is somewhere to the left of Alice.
s.add(pos("Children", "Bella") < pos("Name", "Alice"))

# 13. Arnold is somewhere to the right of the person who loves fantasy books.
s.add(pos("Name", "Arnold") > pos("BookGenre", "fantasy"))

# 14. The person who loves mystery books is in the fourth house.
s.add(pos("BookGenre", "mystery") == 4)

# 15. The person's child is named Alice is the person who enjoys camping trips.
s.add(pos("Children", "Alice") == pos("Vacation", "camping"))

# 16. The person whose mother's name is Kailyn is the person who likes going on cruises.
s.add(pos("Mother", "Kailyn") == pos("Vacation", "cruise"))

# 17. There are two houses between the person who loves fantasy books and The person whose mother's name is Aniya.
s.add(Abs(pos("BookGenre", "fantasy") - pos("Mother", "Aniya")) == 3)

# 18. The person who loves fantasy books is Carol.
s.add(pos("BookGenre", "fantasy") == pos("Name", "Carol"))

# 19. The person who likes going on cruises is the person who loves biography books.
s.add(pos("Vacation", "cruise") == pos("BookGenre", "biography"))

# 20. The person who loves fantasy books is in the third house.
s.add(pos("BookGenre", "fantasy") == 3)

# 21. The person whose mother's name is Aniya is the person who loves romance books.
s.add(pos("Mother", "Aniya") == pos("BookGenre", "romance"))

# 22. The person whose mother's name is Janelle is not in the fourth house.
s.add(pos("Mother", "Janelle") != 4)

# 23. The person's child is named Fred is not in the fourth house.
s.add(pos("Children", "Fred") != 4)

# 24. The person who loves biography books is not in the second house.
s.add(pos("BookGenre", "biography") != 2)

# 25. There are two houses between The person whose mother's name is Holly and Eric.
s.add(Abs(pos("Mother", "Holly") - pos("Name", "Eric")) == 3)

# Solve the puzzle
if s.check() == sat:
    m = s.model()
    rows = []
    header = ["House"] + [cat_name for cat_name, _ in categories]
    for position in range(1, NUM_POSITIONS + 1):
        row = [str(position)]
        for cat_idx, (cat_name, item_list) in enumerate(categories):
            for item_idx, item_str in enumerate(item_list):
                if m.evaluate(Vars[cat_idx][item_idx]).as_long() == position:
                    row.append(item_str)
                    break
        rows.append(row)
    result_dict = {"header": header, "rows": rows}

    cnt = 1
    block = []
    for cat_idx, (cat_name, item_list) in enumerate(categories):
        for i in range(NUM_POSITIONS):
            block.append(Vars[cat_idx][i] != m[Vars[cat_idx][i]])
    s.add(Or(block))
    while s.check() == sat:
        m = s.model()
        cnt += 1
        block = []
        for cat_idx, (cat_name, item_list) in enumerate(categories):
            for i in range(NUM_POSITIONS):
                block.append(Vars[cat_idx][i] != m[Vars[cat_idx][i]])
        s.add(Or(block))
    print(f"sat:{cnt}")
else:
    print(f"error")