a

iq_mini_4489/src/solver.py

@@ -0,0 +1,121 @@
+# Given the position of the pegs, we solve the puzzle
+
+from board import Board, generate_orientations
+from typing import Tuple, List
+import pieces
+import numpy as np
+import colorama
+
+def backtrack(board: Board, pieces_list: List[str], piece_index: int,
+              orientations: dict, solutions: List[np.ndarray]) -> None:
+    """
+    Backtracking to place each piece in the board in all possible ways.
+    """
+    
+    # if full, check if solution
+    if piece_index == len(pieces_list):
+        if board.is_full():
+            solutions.append(board.grid.copy())
+        return
+
+    piece_name = pieces_list[piece_index]
+    all_orientations = orientations[piece_name]
+
+    for orientation in all_orientations:
+        rows, cols = orientation.shape
+        for row in range(board.rows - rows + 1):
+            for col in range(board.cols - cols + 1):
+                if board.placeable(orientation, (row, col)):
+                    board.place(orientation, (row, col), piece_index + 1)
+                    backtrack(board, pieces_list, piece_index + 1, orientations, solutions)
+                    board.remove(orientation, (row, col))
+
+
+
+def solve(xPin: Tuple[int, int], yPin: Tuple[int, int], zPin: Tuple[int, int]):
+    """
+    solve the puzzle, avoiding the pins at the specified positions.
+
+    Args:
+        xPin (Tuple[int, int]): row, col of the 'x' pin
+        yPin (Tuple[int, int]): row, col of the 'y' pin
+        zPin (Tuple[int, int]): row, col of the 'z' pin
+    """
+
+    board = Board((5, 5))
+
+    pinned_positions = [xPin, yPin, zPin]
+    for (pr, pc) in pinned_positions:
+        board.grid[pr, pc] = -1 
+
+    orientations_map = {}
+    for piece_name, piece_matrix in pieces.all_pieces.items():
+        orientations_map[piece_name] = generate_orientations(piece_matrix)
+    piece_names = list(pieces.all_pieces.keys())
+    solutions = []
+    backtrack(board, piece_names, 0, orientations_map, solutions)
+    
+    return solutions
+
+
+def colorful_solution(xPin, yPin, zPin) -> str:
+    """
+    Return a colorful representation of the solution.
+    """
+
+    if not (solutions := solve(xPin, yPin, zPin)):
+        return print("No solution found.")
+
+
+    for i, solution in enumerate(solutions):
+        rows, cols = solution.shape
+        result = ""
+        piece_id_to_color = {
+            -1: colorama.Back.BLACK,
+            1: colorama.Back.GREEN,
+            2: colorama.Back.BLUE,
+            3: colorama.Back.YELLOW,
+            4: colorama.Back.MAGENTA,
+            5: colorama.Back.CYAN,
+            6: colorama.Back.RED,
+        }
+        for r in range(rows):
+            for c in range(cols):
+                piece_id = solution[r, c]
+                if piece_id == 0:
+                    result += colorama.Back.LIGHTWHITE_EX + "  "
+                else:
+                    result += piece_id_to_color[piece_id] + "  "
+                
+                result += colorama.Back.RESET
+            result += "\n"
+        print(f"Solution {i+1}")
+        print(result)
+        print("--------------------")
+    print(f"Total solutions: {len(solutions)}")
+
+def try_all_pegs():
+    from tqdm import tqdm    
+    all_solutions = []
+    # Solution entry {x: (row, col), y: (row, col), z: (row, col), solution_count: int}
+    
+    for xPinA in tqdm(range(5)):
+        for yPinA in range(5):
+            for xPinB in range(5):
+                for yPinB in range(5):
+                    for xPinC in range(5):
+                        for yPinC in range(5):
+                            if (xPinA, yPinA) == (xPinB, yPinB) or (xPinA, yPinA) == (xPinC, yPinC) or (xPinB, yPinB) == (xPinC, yPinC):
+                                continue
+                            solutions = solve((xPinA, yPinA), (xPinB, yPinB), (xPinC, yPinC))
+                            if solutions:
+                                all_solutions.append({
+                                    'x': (xPinA, yPinA),
+                                    'y': (xPinB, yPinB),
+                                    'z': (xPinC, yPinC),
+                                    'solution_count': len(solutions)
+                                })
+    return all_solutions
+if __name__ == "__main__":
+    # colorful_solution((0, 2), (2, 0), (3, 3))
+    print(try_all_pegs())