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

puzzle_solver/puzzles/black_box/black_box.py

@@ -0,0 +1,243 @@
+from typing import Optional
+
+from ortools.sat.python import cp_model
+
+from puzzle_solver.core.utils import Pos, get_all_pos, get_pos, get_row_pos, get_col_pos, in_bounds, get_opposite_direction, get_next_pos, Direction
+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, top: list, right: list, bottom: list, left: list, ball_count: Optional[tuple[int, int]] = None, max_travel_steps: Optional[int] = None):
+        assert len(top) == len(bottom), 'top and bottom must be the same length'
+        assert len(left) == len(right), 'left and right must be the same length'
+        self.K = len(top) + len(right) + len(bottom) + len(left)  # side count
+        self.H = len(top)
+        self.V = len(left)
+        if max_travel_steps is None:
+            self.T = self.V * self.H  # maximum travel steps for a beam that bounces an undefined number of times
+        else:
+            self.T = max_travel_steps
+        self.ball_count = ball_count
+        # top and bottom entry cells are at -1 and V
+        self.top_cells = set(get_row_pos(row_idx=-1, H=self.H))
+        self.bottom_cells = set(get_row_pos(row_idx=self.V, H=self.H))
+        # left and right entry cells are at -1 and H
+        self.left_cells = set(get_col_pos(col_idx=-1, V=self.V))
+        self.right_cells = set(get_col_pos(col_idx=self.H, V=self.V))
+
+        self.top_values = top
+        self.right_values = right
+        self.bottom_values = bottom
+        self.left_values = left
+
+        self.model = cp_model.CpModel()
+        self.ball_states: dict[Pos, cp_model.IntVar] = {}
+        # (entry_pos, T, cell_pos, direction) -> True if the beam that entered from the board at "entry_pos" is present in "cell_pos" and is going in the direction "direction" at time T
+        self.beam_states: dict[tuple[Pos, int, Pos, Direction], cp_model.IntVar] = {}
+        # self.beam_states_ending_at: dict[Pos, cp_model.IntVar] = {}
+        self.beam_states_at_t: dict[int, dict[Pos, dict[tuple[Pos, Direction], cp_model.IntVar]]] = {}
+        self.create_vars()
+        print('Total number of variables:', len(self.ball_states), len(self.beam_states))
+        self.add_all_constraints()
+        print('Solving...')
+
+    def get_outside_border(self):
+        top_border = tuple(get_row_pos(row_idx=-1, H=self.H))
+        bottom_border = tuple(get_row_pos(row_idx=self.V, H=self.H))
+        left_border = tuple(get_col_pos(col_idx=-1, V=self.V))
+        right_border = tuple(get_col_pos(col_idx=self.H, V=self.V))
+        return (*top_border, *bottom_border, *left_border, *right_border)
+
+    def get_all_pos_extended(self):
+        return (*self.get_outside_border(), *get_all_pos(self.V, self.H))
+
+    def create_vars(self):
+        for pos in get_all_pos(self.V, self.H):
+            self.ball_states[pos] = self.model.NewBoolVar(f'ball_at:{pos}')
+        for pos in self.get_all_pos_extended():  # NxN board + 4 edges
+            if pos not in self.ball_states:  # pos is not in the board -> its on the edge
+                self.ball_states[pos] = None  # balls can't be on the edge
+
+        for entry_pos in (self.top_cells | self.right_cells | self.bottom_cells | self.left_cells):
+
+            for t in range(self.T):
+                self.beam_states[(entry_pos, t, 'HIT', 'HIT')] = self.model.NewBoolVar(f'beam:{entry_pos}:{t}:HIT:HIT')
+                for cell in self.get_all_pos_extended():
+                    for direction in Direction:
+                        self.beam_states[(entry_pos, t, cell, direction)] = self.model.NewBoolVar(f'beam:{entry_pos}:{t}:{cell}:{direction}')
+
+        for (entry_pos, t, cell, direction) in self.beam_states.keys():
+            if t not in self.beam_states_at_t:
+                self.beam_states_at_t[t] = {}
+            if entry_pos not in self.beam_states_at_t[t]:
+                self.beam_states_at_t[t][entry_pos] = {}
+            self.beam_states_at_t[t][entry_pos][(cell, direction)] = self.beam_states[(entry_pos, t, cell, direction)]
+
+    def add_all_constraints(self):
+        self.init_beams()
+        self.constrain_beam_movement()
+        self.constrain_final_beam_states()
+        if self.ball_count is not None:
+            s = sum([b for b in self.ball_states.values() if b is not None])
+            b_min, b_max = self.ball_count
+            self.model.Add(s >= b_min)
+            self.model.Add(s <= b_max)
+
+    def init_beams(self):
+        beam_ids = []
+        beam_ids.extend((beam_id, Direction.DOWN) for beam_id in self.top_cells)
+        beam_ids.extend((beam_id, Direction.LEFT) for beam_id in self.right_cells)
+        beam_ids.extend((beam_id, Direction.UP) for beam_id in self.bottom_cells)
+        beam_ids.extend((beam_id, Direction.RIGHT) for beam_id in self.left_cells)
+
+        for (beam_id, direction) in beam_ids:
+            # beam at t=0 is present at beam_id and facing direction
+            self.model.Add(self.beam_states[(beam_id, 0, beam_id, direction)] == 1)
+            for p in self.get_all_pos_extended():
+                for direction in Direction:
+                    if (p, direction) != (beam_id, direction):
+                        self.model.Add(self.beam_states[(beam_id, 0, p, direction)] == 0)
+
+
+    def constrain_beam_movement(self):
+        for t in range(self.T):
+            for entry_pos in self.beam_states_at_t[t].keys():
+                next_state_dict = self.beam_states_at_t[t][entry_pos]
+                self.model.AddExactlyOne(list(next_state_dict.values()))
+                if t == self.T - 1:
+                    continue
+                for (cell, direction), prev_state in next_state_dict.items():
+                    self.constrain_next_beam_state(entry_pos, t+1, cell, direction, prev_state)
+
+
+    def constrain_next_beam_state(self, entry_pos: Pos, t: int, cur_pos: Pos, direction: Direction, prev_state: cp_model.IntVar):
+        if cur_pos == "HIT":  # a beam that was "HIT" stays "HIT"
+            self.model.Add(self.beam_states[(entry_pos, t, "HIT", "HIT")] == 1).OnlyEnforceIf(prev_state)
+            return
+        pos_ahead = get_next_pos(cur_pos, direction)
+        if not in_bounds(pos_ahead, self.V, self.H) and not in_bounds(cur_pos, self.V, self.H):
+            self.model.Add(self.beam_states[(entry_pos, t, cur_pos, direction)] == 1).OnlyEnforceIf(prev_state)
+            return
+
+        # look at the 3 balls ahead of the beam: thus 8 possible scenarios
+        # A beam with no balls ahead of it -> moves forward in the same direction (1 scenario)
+        # A beam that hits a ball head-on -> beam is "HIT" (4 scenarios)
+        # A beam with a ball in its front-left square and no ball ahead of it -> gets deflected 90 degrees to the right (1 scenario)
+        # A beam with a ball in its front-right square and no ball ahead of it -> gets similarly deflected to the left (1 scenario)
+        # A beam that would in its front-left AND front-right squares -> is reflected (1 scenarios)
+
+        direction_left = {
+            Direction.UP: Direction.LEFT,
+            Direction.LEFT: Direction.DOWN,
+            Direction.DOWN: Direction.RIGHT,
+            Direction.RIGHT: Direction.UP,
+        }[direction]
+        direction_right = {
+            Direction.UP: Direction.RIGHT,
+            Direction.RIGHT: Direction.DOWN,
+            Direction.DOWN: Direction.LEFT,
+            Direction.LEFT: Direction.UP,
+        }[direction]
+        reflected = get_opposite_direction(direction)
+        ball_left_pos = get_next_pos(pos_ahead, direction_left)
+        ball_right_pos = get_next_pos(pos_ahead, direction_right)
+        if in_bounds(pos_ahead, self.V, self.H):
+            ball_ahead = self.ball_states[pos_ahead]
+            ball_ahead_not = ball_ahead.Not()
+        else:
+            ball_ahead = False
+            ball_ahead_not = True
+        if in_bounds(ball_left_pos, self.V, self.H):
+            ball_left = self.ball_states[ball_left_pos]
+            ball_left_not = ball_left.Not()
+        else:
+            ball_left = False
+            ball_left_not = True
+        if in_bounds(ball_right_pos, self.V, self.H):
+            ball_right = self.ball_states[ball_right_pos]
+            ball_right_not = ball_right.Not()
+        else:
+            ball_right = False
+            ball_right_not = True
+
+        pos_left = get_next_pos(cur_pos, direction_left)
+        pos_right = get_next_pos(cur_pos, direction_right)
+        pos_reflected = get_next_pos(cur_pos, reflected)
+        if not in_bounds(pos_left, self.V, self.H):
+            pos_left = cur_pos
+        if not in_bounds(pos_right, self.V, self.H):
+            pos_right = cur_pos
+        if not in_bounds(pos_reflected, self.V, self.H):
+            pos_reflected = cur_pos
+
+        # ball head-on -> beam is "HIT"
+        self.model.Add(self.beam_states[(entry_pos, t, "HIT", "HIT")] == 1).OnlyEnforceIf([
+            ball_ahead,
+            prev_state,
+        ])
+        # ball in front-left -> beam is deflected right
+        self.model.Add(self.beam_states[(entry_pos, t, pos_right, direction_right)] == 1).OnlyEnforceIf([
+            ball_ahead_not,
+            ball_left,
+            ball_right_not,
+            prev_state,
+        ])
+        # ball in front-right -> beam is deflected left
+        self.model.Add(self.beam_states[(entry_pos, t, pos_left, direction_left)] == 1).OnlyEnforceIf([
+            ball_ahead_not,
+            ball_left_not,
+            ball_right,
+            prev_state,
+        ])
+        # ball in front-left and front-right -> beam is reflected
+        self.model.Add(self.beam_states[(entry_pos, t, pos_reflected, reflected)] == 1).OnlyEnforceIf([
+            ball_ahead_not,
+            ball_left,
+            ball_right,
+            prev_state,
+        ])
+        # no ball ahead -> beam moves forward in the same direction
+        self.model.Add(self.beam_states[(entry_pos, t, pos_ahead, direction)] == 1).OnlyEnforceIf([
+            ball_ahead_not,
+            ball_left_not,
+            ball_right_not,
+            prev_state,
+        ])
+
+    def constrain_final_beam_states(self):
+        all_values = []
+        all_values.extend([(Pos(x=c, y=-1), top_value) for c, top_value in enumerate(self.top_values)])
+        all_values.extend([(Pos(x=self.H, y=c), right_value) for c, right_value in enumerate(self.right_values)])
+        all_values.extend([(Pos(x=c, y=self.V), bottom_value) for c, bottom_value in enumerate(self.bottom_values)])
+        all_values.extend([(Pos(x=-1, y=c), left_value) for c, left_value in enumerate(self.left_values)])
+        digits = {}
+        hits = []
+        reflects = []
+        for pos, value in all_values:
+            value = str(value)
+            if value.isdecimal():
+                digits.setdefault(value, []).append(pos)
+            elif value == 'H':
+                hits.append(pos)
+            elif value == 'R':
+                reflects.append(pos)
+            else:
+                raise ValueError(f'Invalid value: {value}')
+        for digit, pos_list in digits.items():
+            assert len(pos_list) == 2, f'digit {digit} has {len(pos_list)} positions: {pos_list}'
+            p1, p2 = pos_list
+            self.model.AddExactlyOne([self.beam_states[(p1, self.T-1, p2, direction)] for direction in Direction])
+            self.model.AddExactlyOne([self.beam_states[(p2, self.T-1, p1, direction)] for direction in Direction])
+        for hit in hits:
+            self.model.AddExactlyOne([self.beam_states[(hit, self.T-1, 'HIT', 'HIT')]])
+        for reflect in reflects:
+            self.model.AddExactlyOne([self.beam_states[(reflect, self.T-1, reflect, direction)] for direction in Direction])
+
+    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(self.ball_states[pos]) for pos in get_all_pos(self.V, self.H)})
+        def callback(single_res: SingleSolution):
+            print("Solution found")
+            print(combined_function(self.V, self.H, center_char=lambda r, c: 'O' if single_res.assignment[get_pos(x=c, y=r)] else ''))
+        generic_solve_all(self, board_to_solution, callback=callback if verbose else None, verbose=verbose)

puzzle_solver/puzzles/branches/branches.py

@@ -0,0 +1,64 @@
+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, Direction, get_next_pos, get_opposite_direction, set_char, get_ray
+from puzzle_solver.core.utils_ortools import and_constraint, generic_solve_all, SingleSolution
+from puzzle_solver.core.utils_visualizer import combined_function
+
+
+class Board:
+    def __init__(self, board: np.array):
+        assert board.ndim == 2, f'board must be 2d, got {board.ndim}'
+        assert all((c.item() == ' ') or str(c.item()).isdecimal() for c in np.nditer(board)), 'board must contain only space or digits'
+        self.board = board
+        self.V, self.H = board.shape
+        self.model = cp_model.CpModel()
+        self.model_vars: dict[tuple[Pos, Direction], cp_model.IntVar] = {}
+        self.create_vars()
+        self.add_all_constraints()
+
+    def create_vars(self):
+        for pos in get_all_pos(self.V, self.H):
+            for direction in Direction:
+                self.model_vars[(pos, direction)] = self.model.NewBoolVar(f'{pos}:{direction}')
+
+    def add_all_constraints(self):
+        for pos in get_all_pos(self.V, self.H):
+            dir_vars = [self.model_vars[(pos, direction)] for direction in Direction]
+            c = get_char(self.board, pos)
+            if not str(c).isdecimal():
+                self.model.AddAtMostOne(dir_vars)
+                continue
+            self.model.Add(lxp.Sum(dir_vars) == 0)  # spot with number has to be blank
+            self.range_clue(pos, int(c))
+
+    def range_clue(self, pos: Pos, k: int):
+        vis_vars: list[cp_model.IntVar] = []
+        for direction in Direction:  # Build visibility chains in four direction
+            branch_direction = get_opposite_direction(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.model_vars[(p, branch_direction)] for p in ray[:idx+1]])
+                vis_vars.append(v)
+        self.model.Add(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:
+            return SingleSolution(assignment={pos: direction.name[0] for (pos, direction), var in board.model_vars.items() if solver.Value(var) == 1})
+        def callback(single_res: SingleSolution):
+            print("Solution found")
+            d = {'U': Direction.UP, 'D': Direction.DOWN, 'L': Direction.LEFT, 'R': Direction.RIGHT}
+            opp_d_char = {'U': 'D', 'D': 'U', 'L': 'R', 'R': 'L'}
+            arr = np.full((self.V, self.H), '', dtype=object)
+            center_char = np.full((self.V, self.H), '', dtype=object)
+            for pos, direction in single_res.assignment.items():
+                opp_direction = get_opposite_direction(d[direction])
+                if single_res.assignment.get(get_next_pos(pos, opp_direction), '') == direction:
+                    set_char(arr, pos, direction + opp_d_char[direction])
+                else:
+                    set_char(arr, pos, direction)
+                    set_char(center_char, pos, '*')
+            print(combined_function(self.V, self.H, center_char=lambda r, c: str(self.board[r, c]).strip() or center_char[r, c].strip(), special_content=lambda r, c: arr[r, c]))
+        return generic_solve_all(self, board_to_solution, callback=callback if verbose else None, verbose=verbose)

puzzle_solver/puzzles/bridges/bridges.py

@@ -0,0 +1,104 @@
+from collections import defaultdict
+
+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_row_pos, get_col_pos
+from puzzle_solver.core.utils_ortools import generic_solve_all, SingleSolution
+
+
+class Board:
+    def __init__(self, board: np.array, max_bridges_per_direction: int = 2):
+        assert board.ndim == 2, f'board must be 2d, got {board.ndim}'
+        assert all(c.item() == ' ' or str(c.item()).isdecimal() for c in np.nditer(board)), 'board must contain only spaces or digits'
+        self.board = board
+        self.V = board.shape[0]
+        self.H = board.shape[1]
+        self.max_bridges_per_direction = max_bridges_per_direction
+        self.horiz_bridges: set[tuple[Pos, Pos]] = set()
+        self.vert_bridges: set[tuple[Pos, Pos]] = set()
+
+        self.model = cp_model.CpModel()
+        self.model_vars: dict[tuple[Pos, Pos], cp_model.IntVar] = {}
+        self.is_bridge_active: dict[tuple[Pos, Pos], cp_model.IntVar] = {}
+
+        self.init_bridges()
+        self.create_vars()
+        self.add_all_constraints()
+
+    def init_bridges(self):
+        for row_i in range(self.V):
+            cells_in_row = [i for i in get_row_pos(row_i, self.H) if get_char(self.board, i) != ' ']
+            for cell_i in range(len(cells_in_row) - 1):
+                self.horiz_bridges.add((cells_in_row[cell_i], cells_in_row[cell_i + 1]))
+        for col_i in range(self.H):
+            cells_in_col = [i for i in get_col_pos(col_i, self.V) if get_char(self.board, i) != ' ']
+            for cell_i in range(len(cells_in_col) - 1):
+                self.vert_bridges.add((cells_in_col[cell_i], cells_in_col[cell_i + 1]))
+
+    def create_vars(self):
+        for bridge in self.horiz_bridges | self.vert_bridges:
+            self.model_vars[bridge] = self.model.NewIntVar(0, self.max_bridges_per_direction, f'{bridge}')
+            self.is_bridge_active[bridge] = self.model.NewBoolVar(f'{bridge}:is_active')
+            self.model.Add(self.model_vars[bridge] == 0).OnlyEnforceIf(self.is_bridge_active[bridge].Not())
+            self.model.Add(self.model_vars[bridge] > 0).OnlyEnforceIf(self.is_bridge_active[bridge])
+
+    def add_all_constraints(self):
+        self.constrain_sums()
+        self.constrain_no_overlapping_bridges()
+
+    def constrain_sums(self):
+        for pos in get_all_pos(self.V, self.H):
+            c = get_char(self.board, pos)
+            if c == ' ':
+                continue
+            all_pos_bridges = [bridge for bridge in self.horiz_bridges if pos in bridge] + [bridge for bridge in self.vert_bridges if pos in bridge]
+            self.model.Add(lxp.sum([self.model_vars[bridge] for bridge in all_pos_bridges]) == int(c))
+
+    def constrain_no_overlapping_bridges(self):
+        for horiz_bridge in self.horiz_bridges:
+            for vert_bridge in self.vert_bridges:
+                if self.is_overlapping(horiz_bridge, vert_bridge):
+                    self.model.AddImplication(self.is_bridge_active[horiz_bridge], self.is_bridge_active[vert_bridge].Not())
+                    self.model.AddImplication(self.is_bridge_active[vert_bridge], self.is_bridge_active[horiz_bridge].Not())
+
+    def is_overlapping(self, horiz_bridge: tuple[Pos, Pos], vert_bridge: tuple[Pos, Pos]) -> bool:
+        assert vert_bridge[0].x == vert_bridge[1].x, 'vertical bridge must have constant x'
+        assert horiz_bridge[0].y == horiz_bridge[1].y, 'horizontal bridge must have constant y'
+        xvert = vert_bridge[0].x
+        yvert_min = min(vert_bridge[0].y, vert_bridge[1].y)
+        yvert_max = max(vert_bridge[0].y, vert_bridge[1].y)
+
+        xhoriz_min = min(horiz_bridge[0].x, horiz_bridge[1].x)
+        xhoriz_max = max(horiz_bridge[0].x, horiz_bridge[1].x)
+        yhoriz = horiz_bridge[0].y
+
+        # no equals because that's what the puzzle says
+        x_contained = xhoriz_min < xvert < xhoriz_max
+        y_contained = yvert_min < yhoriz < yvert_max
+        return x_contained and y_contained
+
+    def solve_and_print(self, verbose: bool = True):
+        def board_to_solution(board: Board, solver: cp_model.CpSolverSolutionCallback) -> SingleSolution:
+            assignment = defaultdict(lambda: [0, 0, 0, 0])
+            for bridge in board.horiz_bridges:
+                v = solver.Value(board.model_vars[bridge])
+                assignment[bridge[0]][0] += v
+                assignment[bridge[1]][1] += v
+            for bridge in board.vert_bridges:
+                v = solver.Value(board.model_vars[bridge])
+                assignment[bridge[0]][2] += v
+                assignment[bridge[1]][3] += v
+            # convert to tuples
+            assignment = {pos: tuple(assignment[pos]) for pos in assignment.keys()}
+            return SingleSolution(assignment=assignment)
+        def callback(single_res: SingleSolution):
+            print("Solution found")
+            res = np.full((self.V, self.H), '    ', dtype=object)
+            for pos, (h1, h2, v1, v2) in single_res.assignment.items():
+                c = str(h1) + str(h2) + str(v1) + str(v2)
+                set_char(res, pos, c)
+            for row in res:
+                print('|' + '|'.join(row) + '|\n')
+        return generic_solve_all(self, board_to_solution, callback=callback if verbose else None, verbose=verbose, max_solutions=20)

puzzle_solver/puzzles/chess_range/chess_melee.py

@@ -0,0 +1,6 @@
+from .chess_range import Board as RangeBoard
+
+class Board(RangeBoard):
+    def __init__(self, pieces: list[str], colors: list[str]):
+        super().__init__(pieces=pieces, colors=colors)
+