1010from collections import defaultdict
1111from typing import Dict , List , Tuple , Optional , Set
1212from importlib import resources
13- from z3 import Solver , And , Or , Not , Bool , PbEq
13+ from z3 import Solver , And , Or , Not , Bool , PbEq , sat
1414
1515
1616# ============== ZebraLogicProblem ==============
@@ -32,6 +32,8 @@ def __init__(self, problem: dict):
3232 self .houses = list (range (self .n_houses ))
3333 self .features = self ._make_features (problem ['features' ])
3434 self .clue_texts = problem ['clues' ]
35+ for clue_ir in problem ['clue_irs' ]:
36+ self .apply_ir (clue_ir )
3537
3638 def _make_features (self , features_domains : dict ) -> dict :
3739 """Create Z3 boolean variables for each feature/value/house combination.
@@ -73,7 +75,6 @@ def compile_constraint(self, ir: dict):
7375 same_house: Two entities at same position
7476 pos_relation: Spatial relationship between entities
7577 """
76- from z3 import And , Or , Not
7778
7879 t = ir ["type" ]
7980
@@ -106,11 +107,15 @@ def B(h): return self.features[f2][v2][h]
106107
107108 def k_to_left (k ):
108109 clauses = []
110+
111+ # A is somewhere to the left of B
109112 if k == "?" :
110113 for h1 in self .houses :
111114 for h2 in self .houses :
112115 if h1 < h2 :
113116 clauses .append (And (A (h1 ), B (h2 )))
117+
118+ # A at h, B at h + k
114119 else :
115120 k = int (k )
116121 for h in range (0 , self .n_houses - k ):
@@ -119,11 +124,15 @@ def k_to_left(k):
119124
120125 def k_to_right (k ):
121126 clauses = []
127+
128+ # A is somewhere to the right of B
122129 if k == "?" :
123130 for h1 in self .houses :
124131 for h2 in self .houses :
125132 if h1 > h2 :
126133 clauses .append (And (A (h1 ), B (h2 )))
134+
135+ # A at h, B at h - k
127136 else :
128137 k = int (k )
129138 for h in range (k , self .n_houses ):
@@ -154,12 +163,15 @@ def apply_ir(self, ir):
154163 compiled_disjuncts = []
155164 for disjunct in ir :
156165 assert isinstance (disjunct , list ), "Each disjunct must be a list of conjuncts"
157- compiled_conjuncts = [self .compile_constraint (c ) for c in disjunct ]
166+ compiled_conjuncts = []
167+ for conjunct in disjunct :
168+ assert isinstance (conjunct , dict ), "Each conjunct must be a dict representing an IR"
169+ compiled = self .compile_constraint (conjunct )
170+ compiled_conjuncts .append (compiled )
158171 compiled_disjuncts .append (And (* compiled_conjuncts ))
159172 self .solver .add (Or (* compiled_disjuncts ))
160173
161174 @property
162175 def is_satisfiable (self ) -> bool :
163176 """Check if the current constraints are satisfiable."""
164- from z3 import sat
165177 return self .solver .check () == sat
0 commit comments