multi-puzzle-solver 1.1.8__py3-none-any.whl

puzzle_solver/puzzles/aquarium/aquarium.py

@@ -0,0 +1,97 @@
+import numpy as np
+from ortools.sat.python import cp_model
+
+from puzzle_solver.core.utils import Pos, get_all_pos, get_char, get_neighbors4, get_row_pos, get_col_pos, get_pos
+from puzzle_solver.core.utils_ortools import generic_solve_all, SingleSolution
+from puzzle_solver.core.utils_visualizer import combined_function, id_board_to_wall_fn
+
+
+def _sanity_check(board: np.array):  # percolation check
+    V, H = board.shape
+    visited: set[Pos] = set()
+    finished_islands: set[int] = set()
+    def dfs(pos: Pos, target_i: int):
+        if pos in visited:
+            return
+        visited.add(pos)
+        for neighbor in get_neighbors4(pos, V, H):
+            if neighbor in visited:
+                continue
+            neighbor_i = int(get_char(board, neighbor))
+            if neighbor_i == target_i:
+                dfs(neighbor, target_i)
+    for pos in get_all_pos(V, H):
+        if pos in visited:
+            continue
+        current_i = int(get_char(board, pos))
+        assert current_i not in finished_islands, f'island {current_i} already finished'
+        dfs(pos, current_i)
+        finished_islands.add(current_i)
+    assert len(finished_islands) == len(set(board.flatten())), 'board is not connected'
+
+class Board:
+    def __init__(self, board: np.array, top: np.array, side: np.array):
+        assert board.ndim == 2, f'board must be 2d, got {board.ndim}'
+        _sanity_check(board)
+        self.V, self.H = board.shape
+        assert top.ndim == 1 and top.shape[0] == self.H, 'top must be a 1d array of length board width'
+        assert side.ndim == 1 and side.shape[0] == self.V, 'side must be a 1d array of length board height'
+        assert all((str(c.item()).isdecimal() for c in np.nditer(board))), 'board must contain only digits'
+        self.board = board
+        self.top = top
+        self.side = side
+        self.aquarium_numbers = set([int(c.item()) for c in np.nditer(board)])
+        self.aquariums = {i: [pos for pos in get_all_pos(self.V, self.H) if int(get_char(self.board, pos)) == i] for i in self.aquarium_numbers}
+        self.aquariums_exist_in_row: dict[int, set[int]] = {aq_i: set() for aq_i in self.aquarium_numbers}
+        for aq_i in self.aquarium_numbers:
+            for row in range(self.V):
+                if any(pos.y == row for pos in self.aquariums[aq_i]):
+                    self.aquariums_exist_in_row[aq_i].add(row)
+
+        self.model = cp_model.CpModel()
+        self.model_vars: dict[Pos, cp_model.IntVar] = {}
+        self.is_aquarium_here: dict[tuple[int, int], cp_model.IntVar] = {}  # is the aquarium here?
+
+        self.create_vars()
+        self.add_all_constraints()
+
+    def create_vars(self):
+        for pos in get_all_pos(self.V, self.H):
+            self.model_vars[pos] = self.model.NewBoolVar(f'{pos}')
+        for aq_i in self.aquarium_numbers:
+            for row in self.aquariums_exist_in_row[aq_i]:
+                self.is_aquarium_here[row, aq_i] = self.model.NewBoolVar(f'{row}:{aq_i}')
+
+    def add_all_constraints(self):
+        for aq_i in self.aquarium_numbers:
+            for pos in self.aquariums[aq_i]:
+                self.model.Add(self.is_aquarium_here[pos.y, aq_i] == 1).OnlyEnforceIf(self.model_vars[pos])
+        # aquarium always start from the bottom
+        for aq_i in self.aquarium_numbers:
+            for row in self.aquariums_exist_in_row[aq_i]:
+                if row + 1 not in self.aquariums_exist_in_row[aq_i]:  # (row + 1) is below (row) thus currently (row) is the bottom of the aquarium
+                    continue
+                self.model.Add(self.is_aquarium_here[row + 1, aq_i] == 1).OnlyEnforceIf(self.is_aquarium_here[row, aq_i])
+        for row in range(self.V):
+            for aq_i in self.aquarium_numbers:
+                aq_i_row_pos = [pos for pos in self.aquariums[aq_i] if pos.y == row]
+                for pos in aq_i_row_pos:
+                    # if the aquarium is here, all the squares in the row of this aquarium must be filled
+                    self.model.Add(self.model_vars[pos] == 1).OnlyEnforceIf(self.is_aquarium_here[row, aq_i])
+                # if the aquarium is here, at least one square in the row of this aquarium must be filled
+                if len(aq_i_row_pos) > 0:
+                    self.model.Add(sum([self.model_vars[pos] for pos in aq_i_row_pos]) == len(aq_i_row_pos)).OnlyEnforceIf(self.is_aquarium_here[row, aq_i])
+                    self.model.Add(sum([self.model_vars[pos] for pos in aq_i_row_pos]) == 0).OnlyEnforceIf(self.is_aquarium_here[row, aq_i].Not())
+        # force the top and side constraints
+        for col in range(self.H):
+            self.model.Add(sum([self.model_vars[pos] for pos in get_col_pos(col, self.V)]) == self.top[col])
+        for row in range(self.V):
+            self.model.Add(sum([self.model_vars[pos] for pos in get_row_pos(row, self.H)]) == self.side[row])
+
+    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(board.model_vars[pos]) for pos in board.model_vars.keys()})
+        def callback(single_res: SingleSolution):
+            print("Solution found")
+            print(combined_function(self.V, self.H, cell_flags=id_board_to_wall_fn(self.board), center_char=lambda r, c: 'O' if single_res.assignment[get_pos(x=c, y=r)] == 1 else ''))
+        return generic_solve_all(self, board_to_solution, callback=callback if verbose else None, verbose=verbose, max_solutions=99)

puzzle_solver/puzzles/area_51/area_51.py

@@ -0,0 +1,159 @@
+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, get_pos, in_bounds, Direction, get_next_pos, get_neighbors4, get_ray
+from puzzle_solver.core.utils_ortools import and_constraint, force_connected_component, generic_solve_all, SingleSolution
+from puzzle_solver.core.utils_visualizer import combined_function, id_board_to_wall_fn
+
+
+def get_neq_var(model: cp_model.CpModel, a: cp_model.IntVar, b: cp_model.IntVar) -> cp_model.IntVar:
+    eq_var = model.NewBoolVar(f'{a}:{b}:eq')
+    model.Add(a == b).OnlyEnforceIf(eq_var)
+    model.Add(a != b).OnlyEnforceIf(eq_var.Not())
+    return eq_var.Not()
+
+def enforce_groups_opposite_when(model: cp_model.CpModel, group_a: list[cp_model.IntVar], group_b: list[cp_model.IntVar]):
+    gate = model.NewBoolVar(f"gate_opposite_when_b[{group_a}]:{group_b}")
+    a0 = group_a[0]
+    b0 = group_b[0]
+    for v in group_a[1:]:  # all A equal
+        model.Add(v == a0).OnlyEnforceIf(gate)
+    for v in group_b[1:]:  # all B equal
+        model.Add(v == b0).OnlyEnforceIf(gate)
+    model.Add(a0 != b0).OnlyEnforceIf(gate)  # A different from B
+    return gate
+
+class Board:
+    def __init__(self, board: np.array, dots: dict[Pos, str]):
+        assert board.ndim == 2, f'board must be 2d, got {board.ndim}'
+        chars = [c.item().strip() for c in np.nditer(board)]
+        assert all(c in ['', 'A', 'C'] or c.isdecimal() or (c[0] == 'O' and c[1:].isdecimal()) for c in chars), 'board must contain only space or A or C or digits or O followed by a number'
+        self.board = board
+        self.V, self.H = board.shape
+        assert all(in_bounds(pos, self.V+1, self.H+1) and v.strip() in ['B', 'W'] for pos, v in dots.items()), 'dots must be a dictionary of Pos to B or W'
+        self.dots = dots
+
+        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):
+        for pos in get_all_pos(self.V, self.H):
+            self.b[pos] = self.model.NewBoolVar(f"b[{pos}]")
+            self.w[pos] = self.b[pos].Not()
+
+    def add_all_constraints(self):
+        for pos in get_all_pos(self.V, self.H):  # For each numbered cell c with value k
+            k = str(get_char(self.board, pos)).strip()
+            if not k:
+                continue
+            if k[0] == 'O':  # O{number} is a range clue
+                k = int(k[1:])
+                self.range_clues(pos, k)
+            elif k.isdecimal():  # {number} is a number clue
+                sum_white_neighbors = lxp.Sum([self.w[p] for p in get_neighbors4(pos, self.V, self.H)])
+                sum_black_neighbors = 4 - sum_white_neighbors  # cells outside of border are black by default
+                self.model.Add(sum_white_neighbors == int(k)).OnlyEnforceIf(self.b[pos])
+                self.model.Add(sum_black_neighbors == int(k)).OnlyEnforceIf(self.w[pos])
+            elif k == 'A':  # A is alien; must be inside fence
+                self.model.Add(self.w[pos] == 1)
+            elif k == 'C':  # C is cactus; must be outside fence
+                self.model.Add(self.b[pos] == 1)
+        for pos, color in self.dots.items():  # this is the most complex part of the puzzle
+            assert color in ['W', 'B'], f'Invalid color: {color}'
+            if color == 'W':
+                self.add_white_dot_constraints(pos)
+            elif color == 'B':
+                self.add_black_dot_constraints(pos)
+        self.fence_is_single_block()
+
+    def range_clues(self, pos: Pos, k: int):
+        self.model.Add(self.w[pos] == 1)  # Force it white
+        vis_vars: list[cp_model.IntVar] = []
+        for direction in Direction:  # Build visibility chains in four direction
+            ray = get_ray(pos, direction, self.V, self.H)  # cells outward
+            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)
+        self.model.Add(1 + sum(vis_vars) == int(k))  # Sum of visible whites = 1 (itself) + sum(chains) == k
+
+    def get_2_by_2_block_vars(self, pos: Pos) -> tuple[cp_model.IntVar, ...]:
+        # returns: the following 2x2 Y's along with the 8 X's
+        # . X X .
+        # X Y Y X
+        # X Y Y X
+        # . X X .
+        br = pos
+        bl = get_next_pos(br, Direction.LEFT)
+        tr = get_next_pos(br, Direction.UP)
+        tl = get_next_pos(tr, Direction.LEFT)
+        tl_v = self.b.get(tl, self.model.NewConstant(1))
+        tr_v = self.b.get(tr, self.model.NewConstant(1))
+        bl_v = self.b.get(bl, self.model.NewConstant(1))
+        br_v = self.b.get(br, self.model.NewConstant(1))
+        tl_u_v = self.b.get(get_next_pos(tl, Direction.UP), self.model.NewConstant(1))
+        tl_l_v = self.b.get(get_next_pos(tl, Direction.LEFT), self.model.NewConstant(1))
+        tr_u_v = self.b.get(get_next_pos(tr, Direction.UP), self.model.NewConstant(1))
+        tr_r_v = self.b.get(get_next_pos(tr, Direction.RIGHT), self.model.NewConstant(1))
+        bl_l_v = self.b.get(get_next_pos(bl, Direction.LEFT), self.model.NewConstant(1))
+        bl_d_v = self.b.get(get_next_pos(bl, Direction.DOWN), self.model.NewConstant(1))
+        br_d_v = self.b.get(get_next_pos(br, Direction.DOWN), self.model.NewConstant(1))
+        br_r_v = self.b.get(get_next_pos(br, Direction.RIGHT), self.model.NewConstant(1))
+        return tl_v, tr_v, bl_v, br_v, tl_u_v, tl_l_v, tr_u_v, tr_r_v, bl_l_v, bl_d_v, br_d_v, br_r_v
+
+    def add_white_dot_constraints(self, pos: Pos):
+        # for the main 2x2 block, either two horizontal 1x2 rectangles on top of each other or two vertical 1x2 rectangles next to each other
+        tl, tr, bl, br, tl_u, tl_l, tr_u, tr_r, bl_l, bl_d, br_d, br_r = self.get_2_by_2_block_vars(pos)
+        # if the horizontal variant, then at least one of the X's on the left/right must be different
+        # horizontal variant will need at least one of these to be active
+        horiz_different = get_neq_var(self.model, tl, tl_l) + get_neq_var(self.model, tr, tr_r) + get_neq_var(self.model, bl, bl_l) + get_neq_var(self.model, br, br_r)
+        # if the vertical variant, then at least one of the X's on the top/bottom must be different
+        # vertical variant will need at least one of these to be active
+        vert_different = get_neq_var(self.model, tl, tl_u) + get_neq_var(self.model, tr, tr_u) + get_neq_var(self.model, bl, bl_d) + get_neq_var(self.model, br, br_d)
+        horiz_variant_gate = enforce_groups_opposite_when(self.model, [tl, tr], [bl, br])
+        self.model.Add(horiz_different >= 1).OnlyEnforceIf(horiz_variant_gate)
+        vert_variant_gate = enforce_groups_opposite_when(self.model, [tl, bl], [tr, br])
+        self.model.Add(vert_different >= 1).OnlyEnforceIf(vert_variant_gate)
+        self.model.AddBoolOr([horiz_variant_gate, vert_variant_gate])
+
+    def add_black_dot_constraints(self, pos: Pos):
+        # in the 2x2 block, one block is X and the other 3 are ~X
+        tl, tr, bl, br, tl_u, tl_l, tr_u, tr_r, bl_l, bl_d, br_d, br_r = self.get_2_by_2_block_vars(pos)
+        # that one the is X must also have it's 2 corresponding outward neighbors also be X
+        gate_1 = enforce_groups_opposite_when(self.model, [tl, tl_u, tl_l], [tr, bl, br, bl_l, tr_u])  # V1: block tl is the one that is different from the other 3
+        gate_2 = enforce_groups_opposite_when(self.model, [tr, tr_u, tr_r], [tl, bl, br, tl_u, br_r])  # V2: block tr is the one that is different from the other 3
+        gate_3 = enforce_groups_opposite_when(self.model, [bl, bl_l, bl_d], [tl, tr, br, tl_l, br_d])  # V3: block bl is the one that is different from the other 3
+        gate_4 = enforce_groups_opposite_when(self.model, [br, br_d, br_r], [tl, tr, bl, bl_d, tr_r])  # V4: block br is the one that is different from the other 3
+        self.model.AddBoolOr([gate_1, gate_2, gate_3, gate_4])
+
+    def fence_is_single_block(self):
+        # inside the fence, all cells must be connected
+        force_connected_component(self.model, self.w)
+        # outside the fence, all cells must be connected, + outside border is considered black + outside the fence must touch the border otherwise 'outside the fence' is completely enclosed by the fence which is invalid
+        def is_outside_neighbor(p1: Pos, p2: Pos) -> bool:
+            if abs(p1.x - p2.x) + abs(p1.y - p2.y) == 1:  # manhattan distance is 1
+                return True
+            # both are on the border
+            p1_on_border = p1.x == 0 or p1.x == self.H - 1 or p1.y == 0 or p1.y == self.V - 1
+            p2_on_border = p2.x == 0 or p2.x == self.H - 1 or p2.y == 0 or p2.y == self.V - 1
+            return p1_on_border and p2_on_border
+        b_aug = self.b.copy()
+        # add a single fake cell on the outside
+        b_aug[get_pos(x=-1, y=0)] = self.model.NewConstant(1)
+        force_connected_component(self.model, b_aug, is_neighbor=is_outside_neighbor)
+
+    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(board.w[pos]) for pos in get_all_pos(board.V, board.H)})
+        def callback(single_res: SingleSolution):
+            print("Solution:")
+            print(combined_function(self.V, self.H,
+                cell_flags=id_board_to_wall_fn(np.array([[single_res.assignment[get_pos(x=c, y=r)] for c in range(self.H)] for r in range(self.V)]), border_is_wall=False, border_is=1),
+                center_char=lambda r, c: self.board[r, c].strip(),
+            ))
+        return generic_solve_all(self, board_to_solution, callback=callback if verbose else None, verbose=verbose)

puzzle_solver/puzzles/battleships/battleships.py

@@ -0,0 +1,139 @@
+from dataclasses import dataclass
+from collections import defaultdict
+from typing import Optional
+
+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, get_neighbors8, get_row_pos, get_col_pos, get_pos, in_bounds
+from puzzle_solver.core.utils_ortools import generic_solve_all, SingleSolution
+from puzzle_solver.core.utils_visualizer import combined_function
+
+
+@dataclass
+class Ship:
+    is_active: cp_model.IntVar
+    length: int
+    top_left_pos: Pos
+    body: set[Pos]
+    water: set[Pos]
+    mid_body: set[Pos]
+    top_tip: Optional[Pos]
+    bottom_tip: Optional[Pos]
+    left_tip: Optional[Pos]
+    right_tip: Optional[Pos]
+
+class Board:
+    def __init__(self, board: np.array, top: np.array, side: np.array, ship_counts: dict[int, int]):
+        assert board.ndim == 2, f'board must be 2d, got {board.ndim}'
+        self.V, self.H = board.shape
+        assert top.ndim == 1 and top.shape[0] == self.H, 'top must be a 1d array of length board width'
+        assert side.ndim == 1 and side.shape[0] == self.V, 'side must be a 1d array of length board height'
+        assert all((str(c.item()) in [' ', 'W', 'O', 'S', 'U', 'D', 'L', 'R'] for c in np.nditer(board))), 'board must contain only spaces, W, O, S, U, D, L, R'
+        self.board = board
+        self.top = top
+        self.side = side
+        self.ship_counts = ship_counts
+
+        self.model = cp_model.CpModel()
+        self.model_vars: dict[Pos, cp_model.IntVar] = {}
+        self.shipyard: list[Ship] = []  # will contain every possible ship based on ship counts
+
+        self.create_vars()
+        self.init_shipyard()
+        self.add_all_constraints()
+
+    def create_vars(self):
+        for pos in get_all_pos(self.V, self.H):
+            self.model_vars[pos] = self.model.NewBoolVar(f'{pos}:is_ship')
+
+    def get_ship(self, pos: Pos, length: int, orientation: str) -> Optional[Ship]:
+        if length == 1:
+            body = {pos}
+            top_tip = None
+            bottom_tip = None
+            left_tip = None
+            right_tip = None
+        elif orientation == 'horizontal':
+            body = set(get_pos(x=x, y=pos.y) for x in range(pos.x, pos.x + length))
+            top_tip = None
+            bottom_tip = None
+            left_tip = pos
+            right_tip = get_pos(x=pos.x + length - 1, y=pos.y)
+        elif orientation == 'vertical':
+            body = set(get_pos(x=pos.x, y=y) for y in range(pos.y, pos.y + length))
+            left_tip = None
+            right_tip = None
+            top_tip = pos
+            bottom_tip = get_pos(x=pos.x, y=pos.y + length - 1)
+        else:
+            raise ValueError(f'invalid orientation: {orientation}')
+        if any(not in_bounds(p, self.V, self.H) for p in body):
+            return None
+        water = set(p for pos in body for p in get_neighbors8(pos, self.V, self.H)) - body
+        mid_body = body - {top_tip, bottom_tip, left_tip, right_tip} if length > 1 else set()
+        return Ship(
+            is_active=self.model.NewBoolVar(f'{pos}:is_active'), length=length,
+            top_left_pos=pos, body=body, mid_body=mid_body, water=water,
+            top_tip=top_tip, bottom_tip=bottom_tip, left_tip=left_tip, right_tip=right_tip,
+        )
+
+    def init_shipyard(self):
+        for length in self.ship_counts.keys():
+            for pos in get_all_pos(self.V, self.H):
+                for orientation in ['horizontal', 'vertical']:
+                    if length == 1 and orientation == 'vertical':  # prevent double counting 1-length ships
+                        continue
+                    ship = self.get_ship(pos, length, orientation)
+                    if ship is not None:
+                        self.shipyard.append(ship)
+
+    def add_all_constraints(self):
+        # if a ship is active then all its body is active and all its water is inactive
+        pos_to_ships: dict[Pos, list[Ship]] = defaultdict(list)
+        for ship in self.shipyard:
+            for pos in ship.body:
+                self.model.Add(self.model_vars[pos] == 1).OnlyEnforceIf(ship.is_active)
+                pos_to_ships[pos].append(ship)
+            for pos in ship.water:
+                self.model.Add(self.model_vars[pos] == 0).OnlyEnforceIf(ship.is_active)
+        # if a pos is active then exactly one ship can be placed at that position
+        for pos in get_all_pos(self.V, self.H):
+            self.model.Add(lxp.Sum([ship.is_active for ship in pos_to_ships[pos]]) == 1).OnlyEnforceIf(self.model_vars[pos])
+        # force ship counts
+        for length, count in self.ship_counts.items():
+            self.model.Add(lxp.Sum([ship.is_active for ship in self.shipyard if ship.length == length]) == count)
+        # force the initial board placement
+        for pos in get_all_pos(self.V, self.H):
+            c = get_char(self.board, pos)
+            if c == 'S':  # single-length ship
+                self.model.Add(lxp.Sum([ship.is_active for ship in self.shipyard if ship.length == 1 and ship.top_left_pos == pos]) == 1)
+            elif c == 'W':  # water
+                self.model.Add(self.model_vars[pos] == 0)
+            elif c == 'O':  # mid-body of a ship
+                self.model.Add(lxp.Sum([ship.is_active for ship in self.shipyard if pos in ship.mid_body]) == 1)
+            elif c == 'U':  # top tip of a ship
+                self.model.Add(lxp.Sum([ship.is_active for ship in self.shipyard if ship.top_tip == pos]) == 1)
+            elif c == 'D':  # bottom tip of a ship
+                self.model.Add(lxp.Sum([ship.is_active for ship in self.shipyard if ship.bottom_tip == pos]) == 1)
+            elif c == 'L':  # left tip of a ship
+                self.model.Add(lxp.Sum([ship.is_active for ship in self.shipyard if ship.left_tip == pos]) == 1)
+            elif c == 'R':  # right tip of a ship
+                self.model.Add(lxp.Sum([ship.is_active for ship in self.shipyard if ship.right_tip == pos]) == 1)
+            elif c == ' ':  # empty cell
+                pass
+            else:
+                raise ValueError(f'invalid character: {c}')
+        for row in range(self.V):  # force the top counts
+            self.model.Add(lxp.Sum([self.model_vars[p] for p in get_row_pos(row, self.H)]) == self.side[row])
+        for col in range(self.H):  # force the side counts
+            self.model.Add(lxp.Sum([self.model_vars[p] for p in get_col_pos(col, self.V)]) == self.top[col])
+
+    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")
+            print(combined_function(self.V, self.H, is_shaded=lambda r, c: single_res.assignment[get_pos(x=c, y=r)] == 1))
+        return generic_solve_all(self, board_to_solution, callback=callback if verbose else None, verbose=verbose)

puzzle_solver/puzzles/binairo/binairo.py

@@ -0,0 +1,98 @@
+from typing import Optional
+
+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 Direction, Pos, get_all_pos, get_next_pos, get_pos, in_bounds, get_char, get_row_pos, get_col_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, 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
+        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()
+
+    def create_vars(self):
+        for pos in get_all_pos(self.V, self.H):
+            self.model_vars[pos] = self.model.NewBoolVar(f'{pos}')
+
+    def add_all_constraints(self):
+        for pos in get_all_pos(self.V, self.H):  # force clues
+            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)]
+            self.model.Add(lxp.sum(row_vars) == len(row_vars) // 2)
+        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. 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:
+            self.force_unique_double_list([[self.model_vars[pos] for pos in get_row_pos(row, self.H)] for row in range(self.V)])
+            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:
+            self.force_arithmetic(self.arith_rows, Direction.RIGHT, self.V, self.H-1)
+        if self.arith_cols is not None:
+            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 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]]):
+        m = len(model_vars[0])
+        assert m <= 61, f'Too many cells for binary encoding in int64: m={m}, model_vars={model_vars}'
+        codes = []
+        pow2 = [2**k for k in range(m)]
+        for i, line in enumerate(model_vars):
+            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] = {}
+            for pos, var in board.model_vars.items():
+                assignment[pos] = solver.Value(var)
+            return SingleSolution(assignment=assignment)
+        def callback(single_res: SingleSolution):
+            print("Solution found")
+            print(combined_function(self.V, self.H, is_shaded=lambda r, c: single_res.assignment[get_pos(x=c, y=r)] == 1))
+        return generic_solve_all(self, board_to_solution, callback=callback if verbose else None, verbose=verbose)

puzzle_solver/puzzles/binairo/binairo_plus.py

@@ -0,0 +1,7 @@
+import numpy as np
+
+from . import binairo
+
+class Board(binairo.Board):
+    def __init__(self, board: np.array, arith_rows: np.array, arith_cols: np.array):
+        super().__init__(board=board, arith_rows=arith_rows, arith_cols=arith_cols, force_unique=False)