52. N-Queens II

• categories: Code, Interview Question, leetcode, Hard
• source: https://leetcode.com/problems/n-queens-ii
• topics: Backtracking

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.

Idea

Backtracking, place row by row and store used columns, left and right diagonals

class Solution:
    def totalNQueens(self, n: int) -> int:
        ans = 0
        
        def backtrack(columns, left_diagonals, 
                      right_diagonals, row, queens_to_place):
            nonlocal ans
            for col in range(n):
                # Don't use same column and diagonal
                if (col in columns or 
                    row - col in left_diagonals or 
                    row + col in right_diagonals):
                    continue
                
                # Place the queen
                columns.add(col)
                left_diagonals.add(row - col)
                right_diagonals.add(row + col)
                
                if queens_to_place == 1:
                    # Queens are placed 
                    ans += 1  
                else:
                    # Place the next queen
                    backtrack(columns, left_diagonals, 
                              right_diagonals, row + 1, queens_to_place - 1)
                
                # Backtrack
                columns.remove(col)
                left_diagonals.remove(row - col)
                right_diagonals.remove(row + col)
        
        backtrack(set(), set(), set(), 0, n)
        return ans