Fix n-queens problem

backtracking/n_queens.py

@@ -27,21 +27,28 @@ def is_safe(board: list[list[int]], row: int, column: int) -> bool:
 
     >>> is_safe([[0, 0, 0], [0, 0, 0], [0, 0, 0]], 1, 1)
     True
+    >>> is_safe([[0, 1, 0], [0, 0, 0], [0, 0, 0]], 1, 1)
+    False
     >>> is_safe([[1, 0, 0], [0, 0, 0], [0, 0, 0]], 1, 1)
     False
+    >>> is_safe([[0, 0, 1], [0, 0, 0], [0, 0, 0]], 1, 1)
+    False
     """
 
     n = len(board)  # Size of the board
 
-    # Check if there is any queen in the same row, column,
-    # left upper diagonal, and right upper diagonal
+    # Check if there is any queen in the same upper column,
+    # left upper diagonal and right upper diagonal
     return (
-        all(board[i][j] != 1 for i, j in zip(range(row, -1, -1), range(column, n)))
+        all(board[i][j] != 1 for i, j in zip(range(row), [column] * row))
+        and all(
+            board[i][j] != 1
+            for i, j in zip(range(row - 1, -1, -1), range(column - 1, -1, -1))
+        )
         and all(
-            board[i][j] != 1 for i, j in zip(range(row, -1, -1), range(column, -1, -1))
+            board[i][j] != 1
+            for i, j in zip(range(row - 1, -1, -1), range(column + 1, n))
         )
-        and all(board[i][j] != 1 for i, j in zip(range(row, n), range(column, n)))
-        and all(board[i][j] != 1 for i, j in zip(range(row, n), range(column, -1, -1)))
     )