refactor nonogram solver
change
commit 33c2f7d044f53239456399728b7c64d283869bbf
author Alpha Chen <alpha@kejadlen.dev>
date
parent e4e18eed
diff --git a/puzzle-hunt/red-vs-blue.py b/puzzle-hunt/red-vs-blue.py
index c0d5e4c..ee481f9 100644
--- a/puzzle-hunt/red-vs-blue.py
+++ b/puzzle-hunt/red-vs-blue.py
@@ -92,12 +92,7 @@ def row(y):
 def col(x):
   return [cell(x, y) for y in range(0, len(row_rules))]
 
-s = Solver()
-
-for x, col_rule in enumerate(col_rules):
-  stripe_name = f"col-{x}"
-  stripe_constraints = col_rule
-  stripe = col(x)
+def create_constraints(stripe_name, stripe_constraints, stripe, f):
   group_count = len(stripe_constraints)
   group_start = [Int(f"{stripe_name}-{i}-start") for i in range(0, group_count)]
   group_end = [Int(f"{stripe_name}-{i}-end") for i in range(0, group_count)]
@@ -118,33 +113,15 @@ for x, col_rule in enumerate(col_rules):
         And(k >= group_start[i], k <= group_end[i])
         for i in range(0, group_count)
     ]
-    s.add(stripe[k] == Or(*cell_in_specific_group))
+    s.add(f(stripe[k], Or(*cell_in_specific_group)))
 
-for y, row_rule in enumerate(row_rules):
-  stripe_name = f"row-{y}"
-  stripe_constraints = row_rule
-  stripe = row(y)
-  group_count = len(stripe_constraints)
-  group_start = [Int(f"{stripe_name}-{i}-start") for i in range(0, group_count)]
-  group_end = [Int(f"{stripe_name}-{i}-end") for i in range(0, group_count)]
-
-  # # Start and end of each group
-  for i in range(0, group_count):
-    s.add(group_start[i] >= 0, group_start[i] < len(stripe))
-    s.add(group_end[i] >= 0, group_end[i] < len(stripe))
-    s.add(group_end[i] - group_start[i] == stripe_constraints[i] - 1)
+s = Solver()
 
-  # Gap between each group and following group
-  for i, j in pairwise(range(0, group_count)):
-    s.add(group_start[j] >= group_end[i] + 2)
+for x, col_rule in enumerate(col_rules):
+  create_constraints(f"col-{x}", col_rule, col(x), lambda a,b: a == b)
 
-  # Cells
-  for k in range(0, len(stripe)):
-    cell_in_specific_group = [
-        And(k >= group_start[i], k <= group_end[i])
-        for i in range(0, group_count)
-    ]
-    s.add(stripe[k] == Not(Or(*cell_in_specific_group)))
+for y, row_rule in enumerate(row_rules):
+  create_constraints(f"row-{y}", row_rule, row(y), lambda a,b: a != b)
 
 s.check()
 m = s.model()