leetcode backtracking

52. N-Queens II.py

class Solution:
    def totalNQueens(self, n: int) -> int:
        """
        T: O(n!). There's n choices to put the first queen in first row, then not
        more than n(n-2) to put the second one, and etc. Therefore, time complexity is
        n factorial.
        
        We use two 1-D arrays to represetn diagonals because it's a vector.
        """
        self.count = 0
        column = [0 for _ in range(n)] # keep track of columns that contains a queen
        diag1 = [0 for _ in range(2*n-1)] # keep track of diagonals (one direction)
        diag2 = [0 for _ in range(2*n-1)] # keep track of diagonals (another direction)
        
        def backtrack(y = 0):
            if y == n:
                self.count += 1    
            for x in range(n):
                if column[x] or diag1[x+y] or diag2[x-y+n-1]: 
                    # what's the constraints that we stop searching on a certain path: if the place we're going to place queen is right 
                    continue
                # mark any points at this column and their two diagonals
                column[x] = diag1[x+y] = diag2[x-y+n-1] = 1
                backtrack(y+1)
                # backtrack to previous state by removing the position of the previously placed queen and try to proceed again. If that would not work either, backtrack again.
                column[x] = diag1[x+y] = diag2[x-y+n-1] = 0

        backtrack(0)
        return self.count