Update n_queens_math.py

backtracking/n_queens_math.py

@@ -81,68 +81,56 @@
 
 def depth_first_search(
     possible_board: list[int],
-    diagonal_right_collisions: list[int],
-    diagonal_left_collisions: list[int],
+    diagonal_right_collisions: set[int],
+    diagonal_left_collisions: set[int],
     boards: list[list[str]],
     n: int,
 ) -> None:
     """
     >>> boards = []
-    >>> depth_first_search([], [], [], boards, 4)
+    >>> depth_first_search([], set(), set(), boards, 4)
     >>> for board in boards:
     ...     print(board)
-    ['. Q . . ', '. . . Q ', 'Q . . . ', '. . Q . ']
-    ['. . Q . ', 'Q . . . ', '. . . Q ', '. Q . . ']
+    ['. Q . .', '. . . Q', 'Q . . .', '. . Q .']
+    ['. . Q .', 'Q . . .', '. . . Q', '. Q . .']
     """
 
-    # Get next row in the current board (possible_board) to fill it with a queen
     row = len(possible_board)
 
-    # If row is equal to the size of the board it means there are a queen in each row in
-    # the current board (possible_board)
     if row == n:
-        # We convert the variable possible_board that looks like this: [1, 3, 0, 2] to
-        # this: ['. Q . . ', '. . . Q ', 'Q . . . ', '. . Q . ']
         boards.append([". " * i + "Q " + ". " * (n - 1 - i) for i in possible_board])
         return
 
-    # We iterate each column in the row to find all possible results in each row
     for col in range(n):
-        # We apply that we learned previously. First we check that in the current board
-        # (possible_board) there are not other same value because if there is it means
-        # that there are a collision in vertical. Then we apply the two formulas we
-        # learned before:
-        #
-        # 45º: y - x = b or 45: row - col = b
-        # 135º: y + x = b or row + col = b.
-        #
-        # And we verify if the results of this two formulas not exist in their variables
-        # respectively.  (diagonal_right_collisions, diagonal_left_collisions)
-        #
-        # If any or these are True it means there is a collision so we continue to the
-        # next value in the for loop.
         if (
             col in possible_board
-            or row - col in diagonal_right_collisions
-            or row + col in diagonal_left_collisions
+            or (row - col) in diagonal_right_collisions
+            or (row + col) in diagonal_left_collisions
         ):
             continue
 
-        # If it is False we call dfs function again and we update the inputs
+        possible_board.append(col)
+        diagonal_right_collisions.add(row - col)
+        diagonal_left_collisions.add(row + col)
+
         depth_first_search(
-            [*possible_board, col],
-            [*diagonal_right_collisions, row - col],
-            [*diagonal_left_collisions, row + col],
+            possible_board,
+            diagonal_right_collisions,
+            diagonal_left_collisions,
             boards,
             n,
         )
 
+        # Backtracking
+        possible_board.pop()
+        diagonal_right_collisions.remove(row - col)
+        diagonal_left_collisions.remove(row + col)
+
 
 def n_queens_solution(n: int) -> None:
     boards: list[list[str]] = []
-    depth_first_search([], [], [], boards, n)
+    depth_first_search([], set(), set(), boards, n)
 
-    # Print all the boards
     for board in boards:
         for column in board:
             print(column)