multi-puzzle-solver 1.0.7__py3-none-any.whl → 1.0.9__py3-none-any.whl

puzzle_solver/puzzles/binairo/binairo.py

@@ -10,26 +10,21 @@ from puzzle_solver.core.utils_visualizer import combined_function
 
 
 class Board:
-    def __init__(self, board: np.array, arith_rows: Optional[np.array] = None, arith_cols: Optional[np.array] = None, force_unique: bool = True):
+    def __init__(self, board: np.array, arith_rows: Optional[np.array] = None, arith_cols: Optional[np.array] = None, force_unique: bool = True, disallow_three: bool = True):
         assert board.ndim == 2, f'board must be 2d, got {board.ndim}'
+        assert board.shape[0] % 2 == 0 and board.shape[1] % 2 == 0, f'board must have even number of rows and columns, got {board.shape[0]}x{board.shape[1]}'
         assert all(c.item() in [' ', 'B', 'W'] for c in np.nditer(board)), 'board must contain only space or B'
+        assert arith_rows is None or all(isinstance(c.item(), str) and c.item() in [' ', 'x', '='] for c in np.nditer(arith_rows)), 'arith_rows must contain only space, x, or ='
+        assert arith_cols is None or all(isinstance(c.item(), str) and c.item() in [' ', 'x', '='] for c in np.nditer(arith_cols)), 'arith_cols must contain only space, x, or ='
         self.board = board
         self.V, self.H = board.shape
-        if arith_rows is not None:
-            assert arith_rows.ndim == 2, f'arith_rows must be 2d, got {arith_rows.ndim}'
-            assert arith_rows.shape == (self.V, self.H-1), f'arith_rows must be one column less than board, got {arith_rows.shape} for {board.shape}'
-            assert all(isinstance(c.item(), str) and c.item() in [' ', 'x', '='] for c in np.nditer(arith_rows)), 'arith_rows must contain only space, x, or ='
-        if arith_cols is not None:
-            assert arith_cols.ndim == 2, f'arith_cols must be 2d, got {arith_cols.ndim}'
-            assert arith_cols.shape == (self.V-1, self.H), f'arith_cols must be one column and row less than board, got {arith_cols.shape} for {board.shape}'
-            assert all(isinstance(c.item(), str) and c.item() in [' ', 'x', '='] for c in np.nditer(arith_cols)), 'arith_cols must contain only space, x, or ='
         self.arith_rows = arith_rows
         self.arith_cols = arith_cols
         self.force_unique = force_unique
+        self.disallow_three = disallow_three
 
         self.model = cp_model.CpModel()
         self.model_vars: dict[Pos, cp_model.IntVar] = {}
-
         self.create_vars()
         self.add_all_constraints()
 
@@ -39,11 +34,9 @@ class Board:
 
     def add_all_constraints(self):
         for pos in get_all_pos(self.V, self.H):  # force clues
-            c = get_char(self.board, pos)
-            if c == 'B':
-                self.model.Add(self.model_vars[pos] == 1)
-            elif c == 'W':
-                self.model.Add(self.model_vars[pos] == 0)
+            c = get_char(self.board, pos).strip()
+            if c:
+                self.model.Add(self.model_vars[pos] == (c == 'B'))
         # 1. Each row and each column must contain an equal number of white and black circles.
         for row in range(self.V):
             row_vars = [self.model_vars[pos] for pos in get_row_pos(row, self.H)]
@@ -51,69 +44,48 @@ class Board:
         for col in range(self.H):
             col_vars = [self.model_vars[pos] for pos in get_col_pos(col, self.V)]
             self.model.Add(lxp.sum(col_vars) == len(col_vars) // 2)
-        # 2. More than two circles of the same color can't be adjacent.
-        for pos in get_all_pos(self.V, self.H):
-            self.disallow_three_in_a_row(pos, Direction.RIGHT)
-            self.disallow_three_in_a_row(pos, Direction.DOWN)
-
+        # 2. No three consecutive cells of the same color
+        if self.disallow_three:
+            for pos in get_all_pos(self.V, self.H):
+                self.disallow_three_in_a_row(pos, Direction.RIGHT)
+                self.disallow_three_in_a_row(pos, Direction.DOWN)
         # 3. Each row and column is unique.
         if self.force_unique:
-            # a list per row
             self.force_unique_double_list([[self.model_vars[pos] for pos in get_row_pos(row, self.H)] for row in range(self.V)])
-            # a list per column
             self.force_unique_double_list([[self.model_vars[pos] for pos in get_col_pos(col, self.V)] for col in range(self.H)])
-
         # if arithmetic is provided, add constraints for it
         if self.arith_rows is not None:
-            assert self.arith_rows.shape == (self.V, self.H-1), f'arith_rows must be one column less than board, got {self.arith_rows.shape} for {self.board.shape}'
-            for pos in get_all_pos(self.V, self.H-1):
-                c = get_char(self.arith_rows, pos)
-                if c == 'x':
-                    self.model.Add(self.model_vars[pos] != self.model_vars[get_next_pos(pos, Direction.RIGHT)])
-                elif c == '=':
-                    self.model.Add(self.model_vars[pos] == self.model_vars[get_next_pos(pos, Direction.RIGHT)])
+            self.force_arithmetic(self.arith_rows, Direction.RIGHT, self.V, self.H-1)
         if self.arith_cols is not None:
-            assert self.arith_cols.shape == (self.V-1, self.H), f'arith_cols must be one row less than board, got {self.arith_cols.shape} for {self.board.shape}'
-            for pos in get_all_pos(self.V-1, self.H):
-                c = get_char(self.arith_cols, pos)
-                if c == 'x':
-                    self.model.Add(self.model_vars[pos] != self.model_vars[get_next_pos(pos, Direction.DOWN)])
-                elif c == '=':
-                    self.model.Add(self.model_vars[pos] == self.model_vars[get_next_pos(pos, Direction.DOWN)])
-
+            self.force_arithmetic(self.arith_cols, Direction.DOWN, self.V-1, self.H)
 
     def disallow_three_in_a_row(self, p1: Pos, direction: Direction):
         p2 = get_next_pos(p1, direction)
         p3 = get_next_pos(p2, direction)
-        if any(not in_bounds(p, self.V, self.H) for p in [p1, p2, p3]):
-            return
-        self.model.AddBoolOr([
-            self.model_vars[p1],
-            self.model_vars[p2],
-            self.model_vars[p3],
-        ])
-        self.model.AddBoolOr([
-            self.model_vars[p1].Not(),
-            self.model_vars[p2].Not(),
-            self.model_vars[p3].Not(),
-        ])
+        if all(in_bounds(p, self.V, self.H) for p in [p1, p2, p3]):
+            self.model.AddBoolOr([self.model_vars[p1], self.model_vars[p2], self.model_vars[p3]])
+            self.model.AddBoolOr([self.model_vars[p1].Not(), self.model_vars[p2].Not(), self.model_vars[p3].Not()])
 
     def force_unique_double_list(self, model_vars: list[list[cp_model.IntVar]]):
-        if not model_vars or len(model_vars) < 2:
-            return
         m = len(model_vars[0])
-        assert m <= 61, f"Too many cells for binary encoding in int64: m={m}, model_vars={model_vars}"
-
+        assert m <= 61, f'Too many cells for binary encoding in int64: m={m}, model_vars={model_vars}'
         codes = []
-        pow2 = [1 << k for k in range(m)]  # weights for bit positions (LSB at index 0)
+        pow2 = [2**k for k in range(m)]
         for i, line in enumerate(model_vars):
-            code = self.model.NewIntVar(0, (1 << m) - 1, f"code_{i}")
-            # Sum 2^k * r[k] == code
-            self.model.Add(code == sum(pow2[k] * line[k] for k in range(m)))
+            code = self.model.NewIntVar(0, 2**m, f"code_{i}")
+            self.model.Add(code == lxp.weighted_sum(line, pow2))  # Sum 2^k * r[k] == code
             codes.append(code)
-
         self.model.AddAllDifferent(codes)
 
+    def force_arithmetic(self, arith_board: np.array, direction: Direction, V: int, H: int):
+        assert arith_board.shape == (V, H), f'arith_board going {direction} expected shape {V}x{H}, got {arith_board.shape}'
+        for pos in get_all_pos(V, H):
+            c = get_char(arith_board, pos).strip()
+            if c == 'x':
+                self.model.Add(self.model_vars[pos] != self.model_vars[get_next_pos(pos, direction)])
+            elif c == '=':
+                self.model.Add(self.model_vars[pos] == self.model_vars[get_next_pos(pos, direction)])
+
     def solve_and_print(self, verbose: bool = True):
         def board_to_solution(board: Board, solver: cp_model.CpSolverSolutionCallback) -> SingleSolution:
             assignment: dict[Pos, int] = {}

puzzle_solver/puzzles/mathema_grids/mathema_grids.py

@@ -0,0 +1,119 @@
+from dataclasses import dataclass
+import numpy as np
+from ortools.sat.python import cp_model
+from ortools.util.python import sorted_interval_list
+
+from puzzle_solver.core.utils import Direction, Pos, get_char, get_next_pos, get_row_pos, get_col_pos, in_bounds, set_char
+from puzzle_solver.core.utils_ortools import generic_solve_all, SingleSolution
+from puzzle_solver.core.utils_visualizer import combined_function
+
+
+@dataclass
+class var_with_bounds:
+    var: cp_model.IntVar
+    min_value: int
+    max_value: int
+
+
+def _div_bounds(a_min: int, a_max: int, b_min: int, b_max: int) -> tuple[int, int]:
+    assert not (b_min == 0 and b_max == 0), "Denominator interval cannot be [0, 0]."
+    denoms = [b_min, b_max]
+    if 0 in denoms:
+        denoms.remove(0)
+    if b_min <= -1:
+        denoms += [-1]
+    if b_max >= 1:
+        denoms += [1]
+    candidates = [a_min // d for d in denoms] + [a_max // d for d in denoms]
+    return min(candidates), max(candidates)
+
+
+class Board:
+    def __init__(self, board: np.array, digits: list[int]):
+        assert board.ndim == 2, f'board must be 2d, got {board.ndim}'
+        self.board = board
+        self.V, self.H = board.shape
+        assert self.V >= 3 and self.V % 2 == 1, f'board must have at least 3 rows and an odd number of rows. Got {self.V} rows.'
+        assert self.H >= 3 and self.H % 2 == 1, f'board must have at least 3 columns and an odd number of columns. Got {self.H} columns.'
+        self.digits = digits
+        self.domain_values = sorted_interval_list.Domain.FromValues(self.digits)
+        self.domain_values_no_zero = sorted_interval_list.Domain.FromValues([d for d in self.digits if d != 0])
+
+        self.model = cp_model.CpModel()
+        self.model_vars: dict[Pos, cp_model.IntVar] = {}
+        self.create_vars()
+        assert len(self.model_vars) == len(self.digits), f'len(model_vars) != len(digits), {len(self.model_vars)} != {len(self.digits)}'
+        self.model.AddAllDifferent(list(self.model_vars.values()))
+
+    def create_vars(self):
+        for row in range(0, self.V-2, 2):
+            line_pos = [pos for pos in get_row_pos(row, self.H)]
+            self.parse_line(line_pos)
+        for col in range(0, self.H-2, 2):
+            line_pos = [pos for pos in get_col_pos(col, self.V)]
+            self.parse_line(line_pos)
+
+    def parse_line(self, line_pos: list[Pos]) -> list[int]:
+        last_num = get_char(self.board, line_pos[-1])
+        equal_sign = get_char(self.board, line_pos[-2])
+        assert equal_sign == '=', f'last element of line must be =, got {equal_sign}'
+        line_pos = line_pos[:-2]
+        operators = [get_char(self.board, pos) for pos in line_pos[1::2]]
+        assert all(c.strip() in ['+', '-', '*', '/'] for c in operators), f'even indices of line must be operators, got {operators}'
+        digits_pos = line_pos[::2]
+        running_var = self.get_var(digits_pos[0], fixed=get_char(self.board, digits_pos[0]))
+        for pos, operator in zip(digits_pos[1:], operators):
+            running_var = self.apply_operator(operator, running_var, self.get_var(pos, fixed=get_char(self.board, pos)))
+        self.model.Add(running_var.var == int(last_num))
+        return running_var
+
+    def get_var(self, pos: Pos, fixed: str) -> var_with_bounds:
+        if pos not in self.model_vars:
+            domain = self.domain_values_no_zero if self.might_be_denominator(pos) else self.domain_values
+            self.model_vars[pos] = self.model.NewIntVarFromDomain(domain, f'{pos}')
+        if fixed.strip():
+            self.model.Add(self.model_vars[pos] == int(fixed))
+        return var_with_bounds(var=self.model_vars[pos], min_value=min(self.digits), max_value=max(self.digits))
+
+    def might_be_denominator(self, pos: Pos) -> bool:
+        "Important since if the variable might be a denominator and the domain includes 0 then ortools immediately sets the model as INVALID"
+        above_pos = get_next_pos(pos, Direction.UP)
+        left_pos = get_next_pos(pos, Direction.LEFT)
+        above_operator = get_char(self.board, above_pos) if in_bounds(above_pos, self.V, self.H) else None
+        left_operator = get_char(self.board, left_pos) if in_bounds(left_pos, self.V, self.H) else None
+        return above_operator == '/' or left_operator == '/'
+
+    def apply_operator(self, operator: str, a: var_with_bounds, b: var_with_bounds) -> var_with_bounds:
+        assert operator in ['+', '-', '*', '/'], f'invalid operator: {operator}'
+        if operator == "+":
+            lo = a.min_value + b.min_value
+            hi = a.max_value + b.max_value
+            res = self.model.NewIntVar(lo, hi, "sum")
+            self.model.Add(res == a.var + b.var)
+        elif operator == "-":
+            lo = a.min_value - b.max_value
+            hi = a.max_value - b.min_value
+            res = self.model.NewIntVar(lo, hi, "diff")
+            self.model.Add(res == a.var - b.var)
+        elif operator == "*":
+            cands = [a.min_value*b.min_value, a.min_value*b.max_value, a.max_value*b.min_value, a.max_value*b.max_value]
+            lo, hi = min(cands), max(cands)
+            res = self.model.NewIntVar(lo, hi, "prod")
+            self.model.AddMultiplicationEquality(res, [a.var, b.var])
+        elif operator == "/":
+            self.model.Add(b.var != 0)
+            lo, hi = _div_bounds(a.min_value, a.max_value, b.min_value, b.max_value)
+            res = self.model.NewIntVar(lo, hi, "quot")
+            self.model.AddDivisionEquality(res, a.var, b.var)
+        return var_with_bounds(res, lo, hi)
+
+    def solve_and_print(self, verbose: bool = True):
+        def board_to_solution(board: Board, solver: cp_model.CpSolverSolutionCallback) -> SingleSolution:
+            return SingleSolution(assignment={pos: solver.Value(var) for pos, var in board.model_vars.items()})
+        def callback(single_res: SingleSolution):
+            print("Solution found")
+            output_board = self.board.copy()
+            for pos, var in single_res.assignment.items():
+                set_char(output_board, pos, str(var))
+            print(combined_function(self.V, self.H, center_char=lambda r, c: str(output_board[r, c])))
+        return generic_solve_all(self, board_to_solution, callback=callback if verbose else None, verbose=verbose)