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")