The n-queens puzzle is the problem of placing n queens on an n x n chessboard such that no two queens attack each other.
Given an integer n, return the number of distinct solutions to the n-queens puzzle.
class Solution: def totalNQueens(self, n: int) -> int: def genRestricted(restricted, r, c): restricted = set(restricted) for row in range(n): restricted.add((row, c)) for col in range(n): restricted.add((r, col)) movements = [[-1, -1], [-1, 1], [1, -1], [1, 1]] for movement in movements: row, col = r, c while 0 <= row < n and 0 <= col < n: restricted.add((row, col)) row += movement col += movement return restricted def gen(row, col, curCount, restricted): count, total_count = curCount, 0 for r in range(row, n): for c in range(col, n): if (r, c) not in restricted: count += 1 if count == n: total_count += 1 total_count += gen(row + 1, 0, count, genRestricted(restricted, r, c)) count -= 1 return total_count return gen(0, 0, 0, set())
It fails at n=8. I can’t figure out why, and how to have less iterations. It seems I am already doing the minimum iterations possible.
restricted set seems wasteful, both time- and space-wise. At the end of the successful recursion,
n levels deep it grows to
n^2 size, which drives the total complexity to
O(n^3). And it is not really needed. It is much easier to check the availability of the square by looking at the queens already placed (please forgive the chess lingo;
file stand for vertical, and
rank for horizontal):
def square_is_safe(file, rank, queens_placed): for queen_rank, queen_file in enumerate(queens_placed): if queen_file == file: # vertical attack return false if queen_file - file == queen_rank - rank: # diagonal attack return false if queen_file - file == rank - queen_rank: # anti-diagonal attack return false return true
to be used in
def place_queen_at_rank(queens_placed, rank): if rank == n: total_count += 1 return for file in range(0, n): if square_is_safe(file, rank, queens_placed): queens_placed.append(file) place_queen_at_rank(queens_placed, rank + 1) queens_placed.pop()
And there is a plenty of room for the optimization. For example, you may want to special-case the first rank: due to a symmetry, you only need to inspect a half of it (cutting execution time by the factor of 2).