multi-puzzle-solver 1.0.4__py3-none-any.whl → 1.0.6__py3-none-any.whl

puzzle_solver/puzzles/range/range.py

@@ -1,8 +1,9 @@
 import numpy as np
 from ortools.sat.python import cp_model
 
-from puzzle_solver.core.utils import Pos, get_all_pos, set_char, get_neighbors4, in_bounds, Direction, get_next_pos, get_char
+from puzzle_solver.core.utils import Pos, get_all_pos, get_pos, get_neighbors4, in_bounds, Direction, get_next_pos, get_char
 from puzzle_solver.core.utils_ortools import and_constraint, generic_solve_all, SingleSolution, force_connected_component
+from puzzle_solver.core.utils_visualizer import combined_function
 
 
 def get_ray(pos: Pos, V: int, H: int, direction: Direction) -> list[Pos]:
@@ -18,88 +19,51 @@ def get_ray(pos: Pos, V: int, H: int, direction: Direction) -> list[Pos]:
 class Board:
     def __init__(self, clues: np.ndarray):
         assert clues.ndim == 2 and clues.shape[0] > 0 and clues.shape[1] > 0, f'clues must be 2d, got {clues.ndim}'
-        assert all(isinstance(i.item(), int) and i.item() >= -1 for i in np.nditer(clues)), f'clues must be -1 or >= 0, got {list(np.nditer(clues))}'
-        self.V = clues.shape[0]
-        self.H = clues.shape[1]
+        assert all(str(i.item()).strip() == '' or str(i.item()).strip().isdecimal() for i in np.nditer(clues)), f'clues must be empty or a decimal number, got {list(np.nditer(clues))}'
+        self.V, self.H = clues.shape
         self.clues = clues
-        self.model = cp_model.CpModel()
 
-        # Core vars
-        self.b: dict[Pos, cp_model.IntVar] = {}  # 1=black, 0=white
-        self.w: dict[Pos, cp_model.IntVar] = {}  # 1=white, 0=black
+        self.model = cp_model.CpModel()
+        self.b: dict[Pos, cp_model.IntVar] = {}
+        self.w: dict[Pos, cp_model.IntVar] = {}
 
         self.create_vars()
         self.add_all_constraints()
 
     def create_vars(self):
-        # Cell color vars
         for pos in get_all_pos(self.V, self.H):
             self.b[pos] = self.model.NewBoolVar(f"b[{pos}]")
-            self.w[pos] = self.model.NewBoolVar(f"w[{pos}]")
-            self.model.AddExactlyOne([self.b[pos], self.w[pos]])
+            self.w[pos] = self.b[pos].Not()
 
     def add_all_constraints(self):
         self.no_adjacent_blacks()
-        self.white_connectivity_percolation()
         self.range_clues()
+        force_connected_component(self.model, self.w)
 
     def no_adjacent_blacks(self):
-        cache = set()
         for p in get_all_pos(self.V, self.H):
             for q in get_neighbors4(p, self.V, self.H):
-                if (p, q) in cache:
-                    continue
-                cache.add((p, q))
                 self.model.Add(self.b[p] + self.b[q] <= 1)
 
-
-    def white_connectivity_percolation(self):
-        force_connected_component(self.model, self.w)
-
     def range_clues(self):
-        # For each numbered cell c with value k:
-        #   - Force it white (cannot be black)
-        #   - Build visibility chains in four directions (excluding the cell itself)
-        #   - Sum of visible whites = 1 (itself) + sum(chains) == k
-        for pos in get_all_pos(self.V, self.H):
-            k = get_char(self.clues, pos)
-            if k == -1:
+        for pos in get_all_pos(self.V, self.H):  # For each numbered cell c with value k
+            k = str(get_char(self.clues, pos)).strip()
+            if not k:
                 continue
-            # Numbered cell must be white
-            self.model.Add(self.b[pos] == 0)
-
-            # Build visibility chains per direction (exclude self)
+            self.model.Add(self.w[pos] == 1)  # Force it white
             vis_vars: list[cp_model.IntVar] = []
-            for direction in Direction:
+            for direction in Direction:  # Build visibility chains in four direction
                 ray = get_ray(pos, self.V, self.H, direction)  # cells outward
-                if not ray:
-                    continue
-                # Chain: v0 = w[ray[0]]; vt = w[ray[t]] & vt-1
-                prev = None
-                for idx, cell in enumerate(ray):
+                for idx in range(len(ray)):
                     v = self.model.NewBoolVar(f"vis[{pos}]->({direction.name})[{idx}]")
+                    and_constraint(self.model, target=v, cs=[self.w[p] for p in ray[:idx+1]])
                     vis_vars.append(v)
-                    if idx == 0:
-                        # v0 == w[cell]
-                        self.model.Add(v == self.w[cell])
-                    else:
-                        and_constraint(self.model, target=v, cs=[self.w[cell], prev])
-                    prev = v
-
-            # 1 (self) + sum(vis_vars) == k
-            self.model.Add(1 + sum(vis_vars) == k)
+            self.model.Add(1 + sum(vis_vars) == int(k))  # Sum of visible whites = 1 (itself) + sum(chains) == k
 
     def solve_and_print(self, verbose: bool = True):
         def board_to_solution(board: Board, solver: cp_model.CpSolverSolutionCallback) -> SingleSolution:
-            assignment: dict[Pos, int] = {}
-            for pos, var in board.b.items():
-                assignment[pos] = solver.Value(var)
-            return SingleSolution(assignment=assignment)
+            return SingleSolution(assignment={pos: solver.Value(board.b[pos]) for pos in get_all_pos(board.V, board.H)})
         def callback(single_res: SingleSolution):
             print("Solution:")
-            res = np.full((self.V, self.H), ' ', dtype=object)
-            for pos in get_all_pos(self.V, self.H):
-                c = 'B' if single_res.assignment[pos] == 1 else ' '
-                set_char(res, pos, c)
-            print(res)
+            print(combined_function(self.V, self.H, is_shaded=lambda r, c: single_res.assignment[get_pos(x=c, y=r)] == 1, center_char=lambda r, c: self.clues[r, c].strip(), text_on_shaded_cells=False))
         return generic_solve_all(self, board_to_solution, callback=callback if verbose else None, verbose=verbose)

puzzle_solver/puzzles/rectangles/rectangles.py

@@ -6,7 +6,7 @@ from ortools.sat.python import cp_model
 
 from puzzle_solver.core.utils import Pos, get_all_pos, in_bounds, set_char, get_char, Direction, get_next_pos
 from puzzle_solver.core.utils_ortools import generic_solve_all, SingleSolution
-from puzzle_solver.core.utils_visualizer import render_grid
+from puzzle_solver.core.utils_visualizer import combined_function
 
 
 def factor_pairs(N: int, upper_limit_i: int, upper_limit_j: int):
@@ -121,6 +121,8 @@ class Board:
                     set_char(res, pos, get_char(res, pos) + 'U')
                 if bottom_pos not in single_res.assignment or single_res.assignment[bottom_pos] != cur:
                     set_char(res, pos, get_char(res, pos) + 'D')
-            print(render_grid(res, center_char=lambda r, c: self.board[r, c] if self.board[r, c] != ' ' else ' '))
+            print(combined_function(self.V, self.H,
+                cell_flags=lambda r, c: res[r, c],
+                center_char=lambda r, c: self.board[r, c] if self.board[r, c] != ' ' else ' '))
 
         return generic_solve_all(self, board_to_solution, callback=callback if verbose else None, verbose=verbose)

puzzle_solver/puzzles/shingoki/shingoki.py

@@ -7,7 +7,7 @@ from ortools.sat.python import cp_model
 
 from puzzle_solver.core.utils import Direction, Pos, get_all_pos, get_neighbors4, get_next_pos, get_char, in_bounds
 from puzzle_solver.core.utils_ortools import generic_solve_all, force_connected_component, and_constraint
-from puzzle_solver.core.utils_visualizer import render_grid
+from puzzle_solver.core.utils_visualizer import combined_function
 
 
 @dataclass(frozen=True)
@@ -154,5 +154,5 @@ class Board:
                     res[min_y][min_x] += 'L'
                 else:
                     raise ValueError(f'Invalid position: {pos} and {neighbor}')
-            print(render_grid(res, center_char='.'))
+            print(combined_function(self.V - 1, self.H - 1, cell_flags=lambda r, c: res[r, c], center_char=lambda r, c: '.'))
         return generic_solve_all(self, board_to_solution, callback=callback if verbose else None, verbose=verbose)

puzzle_solver/puzzles/singles/singles.py

@@ -3,7 +3,7 @@ from ortools.sat.python import cp_model
 
 from puzzle_solver.core.utils import Pos, get_all_pos, get_char, get_neighbors4, get_all_pos_to_idx_dict, get_row_pos, get_col_pos, get_pos
 from puzzle_solver.core.utils_ortools import generic_solve_all, SingleSolution, force_connected_component
-from puzzle_solver.core.utils_visualizer import render_shaded_grid
+from puzzle_solver.core.utils_visualizer import combined_function
 
 
 class Board:
@@ -45,5 +45,9 @@ class Board:
             return SingleSolution(assignment={pos: 1 if solver.Value(val) == 1 else 0 for pos, val in board.B.items()})
         def callback(single_res: SingleSolution):
             print("Solution found")
-            print(render_shaded_grid(self.V, self.H, lambda r, c: single_res.assignment[get_pos(x=c, y=r)] == 1, empty_text=lambda r, c: self.board[r, c]))
+            print(combined_function(self.V, self.H,
+                is_shaded=lambda r, c: single_res.assignment[get_pos(x=c, y=r)] == 1,
+                center_char=lambda r, c: self.board[r, c],
+                text_on_shaded_cells=False
+            ))
         return generic_solve_all(self, board_to_solution, callback=callback if verbose else None, verbose=verbose)

puzzle_solver/puzzles/slant/slant.py

@@ -4,8 +4,9 @@ from dataclasses import dataclass
 import numpy as np
 from ortools.sat.python import cp_model
 
-from puzzle_solver.core.utils import Pos, get_all_pos, set_char, in_bounds, get_pos
+from puzzle_solver.core.utils import Pos, get_all_pos, in_bounds, get_pos
 from puzzle_solver.core.utils_ortools import force_no_loops, generic_solve_all, SingleSolution
+from puzzle_solver.core.utils_visualizer import combined_function
 
 
 @dataclass(frozen=True)
@@ -13,10 +14,10 @@ class Node:
     """The grid is represented as a graph of cells connected to corners."""
     node_type: Union[Literal["Cell"], Literal["Corner"]]
     pos: Pos
-    slant: Union[Literal["//"], Literal["\\"], None]
+    slant: Union[Literal["/"], Literal["\\"], None]
 
     def get_neighbors(self, board_nodes: dict[tuple[str, Pos, Optional[str]], "Node"]) -> list["Node"]:
-        if self.node_type == "Cell" and self.slant == "//":
+        if self.node_type == "Cell" and self.slant == "/":
             n1 = board_nodes[("Corner", get_pos(self.pos.x+1, self.pos.y), None)]
             n2 = board_nodes[("Corner", get_pos(self.pos.x, self.pos.y+1), None)]
             return [n1, n2]
@@ -27,8 +28,8 @@ class Node:
         elif self.node_type == "Corner":
             # 4 cells, 2 cells per slant
             n1 = ("Cell", get_pos(self.pos.x-1, self.pos.y-1), "\\")
-            n2 = ("Cell", get_pos(self.pos.x, self.pos.y-1), "//")
-            n3 = ("Cell", get_pos(self.pos.x-1, self.pos.y), "//")
+            n2 = ("Cell", get_pos(self.pos.x, self.pos.y-1), "/")
+            n3 = ("Cell", get_pos(self.pos.x-1, self.pos.y), "/")
             n4 = ("Cell", get_pos(self.pos.x, self.pos.y), "\\")
             return {board_nodes[n] for n in [n1, n2, n3, n4] if n in board_nodes}
 
@@ -61,9 +62,9 @@ class Board:
 
     def create_vars(self):
         for pos in get_all_pos(self.V, self.H):
-            self.model_vars[(pos, '//')] = self.model.NewBoolVar(f'{pos}://')
+            self.model_vars[(pos, '/')] = self.model.NewBoolVar(f'{pos}:/')
             self.model_vars[(pos, '\\')] = self.model.NewBoolVar(f'{pos}:\\')
-            self.model.AddExactlyOne([self.model_vars[(pos, '//')], self.model_vars[(pos, '\\')]])
+            self.model.AddExactlyOne([self.model_vars[(pos, '/')], self.model_vars[(pos, '\\')]])
         for (pos, slant), v in self.model_vars.items():
             self.nodes[Node(node_type="Cell", pos=pos, slant=slant)] = v
         for pos in get_all_pos(self.V + 1, self.H + 1):
@@ -76,8 +77,8 @@ class Board:
             # when pos is (xi, yi) then it gets a +1 contribution for each:
             # - cell (xi-1, yi-1) is a "\\"
             # - cell (xi, yi) is a "\\"
-            # - cell (xi, yi-1) is a "//"
-            # - cell (xi-1, yi) is a "//"
+            # - cell (xi, yi-1) is a "/"
+            # - cell (xi-1, yi) is a "/"
             xi, yi = pos.x, pos.y
             tl_pos = get_pos(xi-1, yi-1)
             br_pos = get_pos(xi, yi)
@@ -85,8 +86,8 @@ class Board:
             bl_pos = get_pos(xi-1, yi)
             tl_var = self.model_vars[(tl_pos, '\\')] if in_bounds(tl_pos, self.V, self.H) else 0
             br_var = self.model_vars[(br_pos, '\\')] if in_bounds(br_pos, self.V, self.H) else 0
-            tr_var = self.model_vars[(tr_pos, '//')] if in_bounds(tr_pos, self.V, self.H) else 0
-            bl_var = self.model_vars[(bl_pos, '//')] if in_bounds(bl_pos, self.V, self.H) else 0
+            tr_var = self.model_vars[(tr_pos, '/')] if in_bounds(tr_pos, self.V, self.H) else 0
+            bl_var = self.model_vars[(bl_pos, '/')] if in_bounds(bl_pos, self.V, self.H) else 0
             self.model.Add(sum([tl_var, tr_var, bl_var, br_var]) == number)
         board_nodes = {(node.node_type, node.pos, node.slant): node for node in self.nodes.keys()}
         self.neighbor_dict = {node: node.get_neighbors(board_nodes) for node in self.nodes.keys()}
@@ -106,12 +107,5 @@ class Board:
             return SingleSolution(assignment=assignment)
         def callback(single_res: SingleSolution):
             print("Solution found")
-            res = np.full((self.V, self.H), ' ', dtype=object)
-            for pos in get_all_pos(self.V, self.H):
-                set_char(res, pos, '/' if single_res.assignment[pos] == '//' else '\\')
-            print('[')
-            for row in range(self.V):
-                line = '    [ ' + ' '.join(res[row].tolist()) + ' ]'
-                print(line)
-            print(']')
+            print(combined_function(self.V, self.H, center_char=lambda r, c: single_res.assignment[get_pos(x=c, y=r)]))
         return generic_solve_all(self, board_to_solution, callback=callback if verbose else None, verbose=verbose)

puzzle_solver/puzzles/slitherlink/slitherlink.py

@@ -4,7 +4,7 @@ from ortools.sat.python import cp_model
 
 from puzzle_solver.core.utils import Pos, get_all_pos, get_char, set_char, Direction, get_row_pos, get_col_pos, get_next_pos, in_bounds, get_opposite_direction
 from puzzle_solver.core.utils_ortools import generic_solve_all, SingleSolution, force_connected_component
-from puzzle_solver.core.utils_visualizer import render_grid
+from puzzle_solver.core.utils_visualizer import combined_function
 
 
 CellBorder = tuple[Pos, Direction]
@@ -126,5 +126,5 @@ class Board:
                 set_char(res, pos, c)
             # replace " " with "·"
             board = np.where(self.board == ' ', '·', self.board)
-            print(render_grid(cell_flags=res, center_char=board))
+            print(combined_function(self.V, self.H, cell_flags=lambda r, c: res[r, c], center_char=lambda r, c: board[r, c]))
         return generic_solve_all(self, board_to_solution, callback=callback if verbose else None, verbose=verbose, max_solutions=999)

puzzle_solver/puzzles/star_battle/star_battle.py

@@ -3,7 +3,7 @@ from ortools.sat.python import cp_model
 
 from puzzle_solver.core.utils import Pos, get_all_pos, get_char, set_char, get_neighbors8, get_row_pos, get_col_pos
 from puzzle_solver.core.utils_ortools import generic_solve_all, SingleSolution
-from puzzle_solver.core.utils_visualizer import id_board_to_wall_board, render_grid
+from puzzle_solver.core.utils_visualizer import combined_function, id_board_to_wall_fn
 
 
 class Board:
@@ -70,5 +70,8 @@ class Board:
                         set_char(res, pos, ' ')
                     else:
                         set_char(res, pos, '.')
-            print(render_grid(id_board_to_wall_board(self.board), center_char=lambda r, c: res[r][c]))
+            print(combined_function(self.V, self.H,
+                cell_flags=id_board_to_wall_fn(self.board),
+                center_char=lambda r, c: res[r][c]
+            ))
         return generic_solve_all(self, board_to_solution, callback=callback if verbose else None, verbose=verbose)

puzzle_solver/puzzles/stitches/stitches.py

@@ -3,8 +3,9 @@ from typing import Union
 import numpy as np
 from ortools.sat.python import cp_model
 
-from puzzle_solver.core.utils import Pos, get_all_pos, get_char, set_char, get_next_pos, Direction, get_row_pos, get_col_pos, in_bounds, get_opposite_direction
+from puzzle_solver.core.utils import Pos, get_all_pos, get_char, get_next_pos, Direction, get_row_pos, get_col_pos, in_bounds, get_opposite_direction, get_pos
 from puzzle_solver.core.utils_ortools import generic_solve_all, SingleSolution, and_constraint
+from puzzle_solver.core.utils_visualizer import combined_function, id_board_to_wall_fn
 
 
 class Board:
@@ -40,10 +41,6 @@ class Board:
                     self.block_neighbors.setdefault((block_i, block_j), set()).add((pos, direction, neighbor, opposite_direction))
                     self.valid_stitches.add((pos, neighbor))
                     self.valid_stitches.add((neighbor, pos))
-        # for pair in self.block_neighbors.keys():
-        #     print(pair, self.block_neighbors[pair])
-        # print('top empties', self.top_empties)
-        # print('side empties', self.side_empties)
 
         self.model = cp_model.CpModel()
         self.model_vars: dict[tuple[Pos, Union[Direction, None]], cp_model.IntVar] = {}
@@ -62,19 +59,15 @@ class Board:
             state = [self.model_vars[(pos, direction)] for direction in Direction]
             state.append(self.model_vars[(pos, None)])
             self.model.AddExactlyOne(state)
-            # print('ONLY 1 DIRECTION. only one', state)
         # If a position points at X (and this is a valid pair) then X has to point at me
         for pos in get_all_pos(self.V, self.H):
             for direction in Direction:
                 neighbor = get_next_pos(pos, direction)
-                if not in_bounds(neighbor, self.V, self.H) or (pos, neighbor) not in self.valid_stitches:
-                    # this is not a valid stitch
+                if not in_bounds(neighbor, self.V, self.H) or (pos, neighbor) not in self.valid_stitches:  # this is not a valid stitch
                     self.model.Add(self.model_vars[(pos, direction)] == 0)
-                    # print(f'Pos {pos} cant be {direction}')
                     continue
                 opposite_direction = get_opposite_direction(direction)
                 self.model.Add(self.model_vars[(pos, direction)] == self.model_vars[(neighbor, opposite_direction)])
-                # print(f'{pos}:{direction} must == {neighbor}:{opposite_direction}')
 
         # all blocks connected exactly N times (N usually 1 but can be 2 or 3)
         for connections in self.block_neighbors.values():
@@ -93,17 +86,11 @@ class Board:
 
     def solve_and_print(self, verbose: bool = True):
         def board_to_solution(board: Board, solver: cp_model.CpSolverSolutionCallback) -> SingleSolution:
-            assignment: dict[Pos, str] = {}
-            for (pos, direction), var in board.model_vars.items():
-                if solver.value(var) == 1:
-                    assignment[pos] = direction.name[0] if direction is not None else ' '
-            return SingleSolution(assignment=assignment)
+            return SingleSolution(assignment={pos: direction.name[0] if direction is not None else ' ' for (pos, direction), var in board.model_vars.items() if solver.Value(var) == 1 and direction is not None})
         def callback(single_res: SingleSolution):
             print("Solution found")
-            res = np.full((self.V, self.H), ' ', dtype=object)
-            for pos in get_all_pos(self.V, self.H):
-                c = get_char(self.board, pos)
-                c = single_res.assignment[pos]
-                set_char(res, pos, c)
-            print(res)
+            print(combined_function(self.V, self.H,
+                cell_flags=id_board_to_wall_fn(self.board),
+                special_content=lambda r, c: single_res.assignment.get(get_pos(x=c, y=r), ''),
+                center_char=lambda r, c: 'O' if get_pos(x=c, y=r) in single_res.assignment else '.'))
         return generic_solve_all(self, board_to_solution, callback=callback if verbose else None, verbose=verbose, max_solutions=9)

puzzle_solver/puzzles/sudoku/sudoku.py

@@ -6,6 +6,7 @@ from ortools.sat.python import cp_model
 
 from puzzle_solver.core.utils import Pos, get_pos, get_all_pos, get_char, set_char, get_row_pos, get_col_pos
 from puzzle_solver.core.utils_ortools import and_constraint, generic_solve_all, or_constraint, SingleSolution
+from puzzle_solver.core.utils_visualizer import combined_function
 
 
 def get_value(board: np.array, pos: Pos) -> Union[int, str]:
@@ -82,7 +83,7 @@ class Board:
             assert block_size is None, 'cannot set block size if blocks are not constrained'
 
         if jigsaw is not None:
-            if self.constrain_blocks is not None:
+            if self.constrain_blocks:
                 print('Warning: jigsaw and blocks are both constrained, are you sure you want to do this?')
             assert jigsaw.ndim == 2, f'jigsaw must be 2d, got {jigsaw.ndim}'
             assert jigsaw.shape[0] == self.V and jigsaw.shape[1] == self.H, 'jigsaw must be the same size as the board'
@@ -186,18 +187,11 @@ class Board:
 
     def solve_and_print(self, verbose: bool = True):
         def board_to_solution(board: Board, solver: cp_model.CpSolverSolutionCallback) -> SingleSolution:
-            assignment: dict[Pos, int] = {}
-            for pos, var in board.model_vars.items():
-                assignment[pos] = solver.value(var)
-            return SingleSolution(assignment=assignment)
+            return SingleSolution(assignment={pos: solver.value(var) for pos, var in board.model_vars.items()})
         def callback(single_res: SingleSolution):
             print("Solution found")
-            res = np.full((self.V, self.H), ' ', dtype=object)
-            for pos in get_all_pos(self.V, self.H):
-                c = get_value(self.board, pos)
-                c = single_res.assignment[pos]
-                set_value(res, pos, c)
-            print(res)
+            val_arr = np.array([[single_res.assignment[get_pos(x=c, y=r)] for c in range(self.H)] for r in range(self.V)])
+            print(combined_function(self.V, self.H, center_char=lambda r, c: val_arr[r, c] if val_arr[r, c] < 10 else chr(val_arr[r, c] - 10 + ord('a'))))
         return generic_solve_all(self, board_to_solution, callback=callback if verbose else None, verbose=verbose)
 
 

puzzle_solver/puzzles/tapa/tapa.py

@@ -6,7 +6,7 @@ from ortools.sat.python import cp_model
 
 from puzzle_solver.core.utils import Direction8, Pos, get_all_pos, set_char, get_char, in_bounds, Direction, get_next_pos, get_pos
 from puzzle_solver.core.utils_ortools import generic_solve_all, SingleSolution, force_connected_component
-from puzzle_solver.core.utils_visualizer import render_shaded_grid
+from puzzle_solver.core.utils_visualizer import combined_function
 
 
 def rotated_assignments_N_nums(Xs: tuple[int, ...], target_length: int = 8) -> set[tuple[bool, ...]]:
@@ -94,5 +94,9 @@ class Board:
                 if len(c) > 3:
                     c = '...'
                 set_char(board_justified, pos, ' ' * (2 - len(c)) + c)
-            print(render_shaded_grid(self.V, self.H, lambda r, c: single_res.assignment[get_pos(x=c, y=r)] == 1, empty_text=lambda r, c: str(board_justified[r, c])))
+            print(combined_function(self.V, self.H,
+                is_shaded=lambda r, c: single_res.assignment[get_pos(x=c, y=r)] == 1,
+                center_char=lambda r, c: str(board_justified[r, c]),
+                text_on_shaded_cells=False
+            ))
         return generic_solve_all(self, board_to_solution, callback=callback if verbose else None, verbose=verbose, max_solutions=5)

puzzle_solver/puzzles/tents/tents.py

@@ -4,107 +4,77 @@ import numpy as np
 from ortools.sat.python import cp_model
 from ortools.sat.python.cp_model import LinearExpr as lxp
 
-from puzzle_solver.core.utils import Pos, get_all_pos, get_char, set_char, get_neighbors8, get_next_pos, Direction, get_row_pos, get_col_pos
+from puzzle_solver.core.utils import Direction, Pos, get_all_pos, get_char, get_neighbors8, get_next_pos, get_row_pos, get_col_pos, get_opposite_direction, get_pos
 from puzzle_solver.core.utils_ortools import generic_solve_all, SingleSolution
+from puzzle_solver.core.utils_visualizer import combined_function
 
 
 class Board:
-    def __init__(self, board: np.array, sides: dict[str, np.array]):
+    def __init__(self, board: np.array, side: np.array, top: np.array):
         assert board.ndim == 2, f'board must be 2d, got {board.ndim}'
-        assert board.shape[0] == board.shape[1], 'board must be square'
-        assert len(sides) == 2, '2 sides must be provided'
-        assert set(sides.keys()) == set(['top', 'side'])
-        assert all(s.ndim == 1 and s.shape[0] == board.shape[0] for s in sides.values()), 'all sides must be equal to board size'
+        assert side.ndim == 1 and side.shape[0] == board.shape[0], 'side must be 1d and equal to board size'
+        assert top.ndim == 1 and top.shape[0] == board.shape[1], 'top must be 1d and equal to board size'
         assert all(c.item() in [' ', 'T'] for c in np.nditer(board)), 'board must contain only space or T'
         self.board = board
-        self.N = board.shape[0]
-        self.star_positions: set[Pos] = {pos for pos in get_all_pos(self.N) if get_char(self.board, pos) == ' '}
-        self.tree_positions: set[Pos] = {pos for pos in get_all_pos(self.N) if get_char(self.board, pos) == 'T'}
+        self.V, self.H = board.shape
+        self.side = side
+        self.top = top
+        self.non_tree_positions: set[Pos] = {pos for pos in get_all_pos(self.V, self.H) if get_char(self.board, pos) == ' '}
+        self.tree_positions: set[Pos] = {pos for pos in get_all_pos(self.V, self.H) if get_char(self.board, pos) == 'T'}
+
         self.model = cp_model.CpModel()
-        self.is_tent = defaultdict(int)
-        self.tent_direction = defaultdict(int)
-        self.sides = sides
+        self.is_tent: dict[Pos, cp_model.IntVar] = defaultdict(int)
+        self.tent_direction: dict[tuple[Pos, Direction], cp_model.IntVar] = defaultdict(int)
         self.create_vars()
         self.add_all_constraints()
 
     def create_vars(self):
-        for pos in self.star_positions:
-            is_tent = self.model.NewBoolVar(f'{pos}:is_tent')
-            tent_direction = self.model.NewIntVar(0, 4, f'{pos}:tent_direction')
-            self.model.Add(tent_direction == 0).OnlyEnforceIf(is_tent.Not())
-            self.model.Add(tent_direction > 0).OnlyEnforceIf(is_tent)
-            self.is_tent[pos] = is_tent
-            self.tent_direction[pos] = tent_direction
+        for pos in self.non_tree_positions:
+            self.is_tent[pos] = self.model.NewBoolVar(f'{pos}:is_tent')
+        for pos in self.tree_positions:
+            for direction in Direction:
+                tent_pos = get_next_pos(pos, direction)
+                if tent_pos not in self.is_tent:
+                    continue
+                opposite_direction = get_opposite_direction(direction)
+                tent_direction = self.model.NewBoolVar(f'{pos}:{direction}')
+                self.model.Add(tent_direction == 0).OnlyEnforceIf(self.is_tent[tent_pos].Not())
+                self.tent_direction[(pos, direction)] = tent_direction
+                self.tent_direction[(tent_pos, opposite_direction)] = tent_direction
 
     def add_all_constraints(self):
         # - There are exactly as many tents as trees.
-        self.model.Add(lxp.sum([self.is_tent[pos] for pos in self.star_positions]) == len(self.tree_positions))
+        self.model.Add(lxp.sum([self.is_tent[pos] for pos in self.non_tree_positions]) == len(self.tree_positions))
         # - no two tents are adjacent horizontally, vertically or diagonally
-        for pos in self.star_positions:
-            for neighbour in get_neighbors8(pos, V=self.N, H=self.N, include_self=False):
-                if get_char(self.board, neighbour) != ' ':
-                    continue
+        for pos in self.non_tree_positions:
+            for neighbour in get_neighbors8(pos, V=self.V, H=self.H, include_self=False):
                 self.model.Add(self.is_tent[neighbour] == 0).OnlyEnforceIf(self.is_tent[pos])
         # - the number of tents in each row and column matches the numbers around the edge of the grid
-        for row in range(self.N):
-            row_vars = [self.is_tent[pos] for pos in get_row_pos(row, self.N)]
-            self.model.Add(lxp.sum(row_vars) == self.sides['side'][row])
-        for col in range(self.N):
-            col_vars = [self.is_tent[pos] for pos in get_col_pos(col, self.N)]
-            self.model.Add(lxp.sum(col_vars) == self.sides['top'][col])
+        for row in range(self.V):
+            if self.side[row] == -1:
+                continue
+            row_vars = [self.is_tent[pos] for pos in get_row_pos(row, H=self.H)]
+            self.model.Add(lxp.sum(row_vars) == self.side[row])
+        for col in range(self.H):
+            if self.top[col] == -1:
+                continue
+            col_vars = [self.is_tent[pos] for pos in get_col_pos(col, V=self.V)]
+            self.model.Add(lxp.sum(col_vars) == self.top[col])
         # - it is possible to match tents to trees so that each tree is orthogonally adjacent to its own tent (but may also be adjacent to other tents).
-        # for each tree, one of the following must be true:
-        # a tent on its left has direction RIGHT
-        # a tent on its right has direction LEFT
-        # a tent on its top has direction DOWN
-        # a tent on its bottom has direction UP
-        for tree in self.tree_positions:
-            self.add_tree_constraints(tree)
-
-    def add_tree_constraints(self, tree_pos: Pos):
-        left_pos = get_next_pos(tree_pos, Direction.LEFT)
-        right_pos = get_next_pos(tree_pos, Direction.RIGHT)
-        top_pos = get_next_pos(tree_pos, Direction.UP)
-        bottom_pos = get_next_pos(tree_pos, Direction.DOWN)
-        var_list = []
-        if left_pos in self.star_positions:
-            aux = self.model.NewBoolVar(f'{tree_pos}:left')
-            self.model.Add(self.tent_direction[left_pos] == Direction.RIGHT.value).OnlyEnforceIf(aux)
-            self.model.Add(self.tent_direction[left_pos] != Direction.RIGHT.value).OnlyEnforceIf(aux.Not())
-            var_list.append(aux)
-        if right_pos in self.star_positions:
-            aux = self.model.NewBoolVar(f'{tree_pos}:right')
-            self.model.Add(self.tent_direction[right_pos] == Direction.LEFT.value).OnlyEnforceIf(aux)
-            self.model.Add(self.tent_direction[right_pos] != Direction.LEFT.value).OnlyEnforceIf(aux.Not())
-            var_list.append(aux)
-        if top_pos in self.star_positions:
-            aux = self.model.NewBoolVar(f'{tree_pos}:top')
-            self.model.Add(self.tent_direction[top_pos] == Direction.DOWN.value).OnlyEnforceIf(aux)
-            self.model.Add(self.tent_direction[top_pos] != Direction.DOWN.value).OnlyEnforceIf(aux.Not())
-            var_list.append(aux)
-        if bottom_pos in self.star_positions:
-            aux = self.model.NewBoolVar(f'{tree_pos}:bottom')
-            self.model.Add(self.tent_direction[bottom_pos] == Direction.UP.value).OnlyEnforceIf(aux)
-            self.model.Add(self.tent_direction[bottom_pos] != Direction.UP.value).OnlyEnforceIf(aux.Not())
-            var_list.append(aux)
-        self.model.AddBoolOr(var_list)
+        # each tent is pointing exactly once at a tree
+        for pos in self.non_tree_positions:
+            var_list = [self.tent_direction[(pos, direction)] for direction in Direction]
+            self.model.Add(lxp.sum(var_list) == 1).OnlyEnforceIf(self.is_tent[pos])
+            self.model.Add(lxp.sum(var_list) == 0).OnlyEnforceIf(self.is_tent[pos].Not())
+        # each tree is pointed at by exactly one tent
+        for pos in self.tree_positions:
+            var_list = [self.tent_direction[(pos, direction)] for direction in Direction]
+            self.model.Add(lxp.sum(var_list) == 1)
 
     def solve_and_print(self, verbose: bool = True):
         def board_to_solution(board: Board, solver: cp_model.CpSolverSolutionCallback) -> SingleSolution:
-            assignment: dict[Pos, int] = {}
-            for pos, var in board.is_tent.items():
-                if isinstance(var, int):
-                    continue
-                assignment[pos] = solver.value(var)
-            return SingleSolution(assignment=assignment)
+            return SingleSolution(assignment={pos: solver.value(var) for pos, var in board.is_tent.items() if not isinstance(var, int)})
         def callback(single_res: SingleSolution):
             print("Solution found")
-            res = np.full((self.N, self.N), ' ', dtype=object)
-            for pos in get_all_pos(self.N):
-                c = get_char(self.board, pos)
-                if c == ' ':
-                    c = single_res.assignment[pos]
-                    c = 'E' if c == 1 else ' '
-                set_char(res, pos, c)
-            print(res)
-        return generic_solve_all(self, board_to_solution, callback=callback if verbose else None, verbose=verbose)
+            print(combined_function(self.V, self.H, center_char=lambda r, c: ('|' if self.board[r][c].strip() else ('▲' if single_res.assignment[get_pos(c, r)] == 1 else ' '))))
+        return generic_solve_all(self, board_to_solution, callback=callback if verbose else None, verbose=verbose, max_solutions=5)