leetcode-backtracking

51. N-Queens.py

class Solution:
    """
    Thought Porcess:
     ----------->x
     0 1 2 3
    0
    1
    2
    3
    
    Idea is when we choose one position in the first row to put the queue, then we recursively check second row, his two diagonal direction and column won't be valid to put queen.
    So, we need three array to keep track where we we can still put quene or not:
        -column
        -diag1
        -diag2
    """
    def solveNQueens(self, n: int) -> List[List[str]]:
        """
        Brute force is to generate all possible ways to put N queens on the 
        board, which is O(n**n), n raised to power of n
        
        With backtracking, Time complexity is O(n!)
        
        """
        ans = []
        column = [0 for _ in range(n)]
        diag1 = [0 for _ in range(2*n-1)] # at most, we have 2n-1 diagonals for each direction
        diag2 = [0 for _ in range(2*n-1)] 
        
        def _add_solution(col_pos):
            solution = []
            for x in col_pos:
                s = ["."] * n
                s[x] = "Q"
                solution.append("".join(s))
            ans.append(solution)
            
        def _valid_place(x, y):
            return False if column[x] or diag1[x+y] or diag2[x-y+n-1] else True
        
        def backtrack(y, col_pos, column, diag1, diag2):
            """
            col_pos to store position of column we can valid put quene
            """
            if y == n:
                # there's no place to put the queen
                _add_solution(col_pos)
            # for each row, we try to place the queen at all possible coloum 
            for x in range(n):
                # if valible, we keep exploring next row
                if _valid_place(x, y):
                    col_pos.append(x) # keep track columns we placeed queen
                    # x-y的變化是-n+1到n-1總共2n-1條, 所以我們x-y+(n-1)
                    column[x] = diag1[x+y] = diag2[x-y+n-1] = 1 # mark columns and their two diagonals could not place queen
                    backtrack(y+1, col_pos, column, diag1, diag2)
                    column[x] = diag1[x+y] = diag2[x-y+n-1] = 0
                    col_pos.pop()

        backtrack(0, [], column, diag1, diag2)
        return ans
    
    def solveNQueens(self, n: int) -> List[List[str]]:
        """
        T:O(n! * n): we take additional constant time to progressively build col_pos
        S:O(n)
        """
        def _avaliable(x, y, col, diag1, diag2):
            """
            if n = 4:
                x-y 在-3到3之間, 我們把它轉到0到2n-1 by adding n-1
            """
            if col[x] or diag1[x+y] or diag2[x-y+n-1]:
                return False
            return True
        
        def _add(x, y, col_pos, col, diag1, diag2):
            col_pos.append(x)
            col[x] = 1
            diag1[x+y] = 1
            diag2[x-y+n-1] = 1
            return 
        
        def _remove(x, y, col_pos, col, diag1, diag2):
            """
            backtrack to previous state
            """
            col_pos.pop()
            col[x] = 0
            diag1[x+y] = 0
            diag2[x-y+n-1] = 0
            return 
        
        def from_col_pos_2_sol(col_pos):
            n = len(col_pos)
            sol = []
            for x in col_pos:
                row = ["."] * n
                row[x] = "Q"
                sol.append("".join(row))
            return sol
            
        def bactrack(y, col_pos, n, col, diag1, diag2):
            """
            y: in y row, we need to put queen, which will be kept track of, in order to 
            deterime if we meet requirement
            col: list of x cooredinate, where we put the queen in this column.
            
            """
            if y == n:
                ans.append(from_col_pos_2_sol(col_pos))
                return
            
            for x in range(n):
                if _avaliable(x,y,col,diag1,diag2):
                    _add(x, y, col_pos, col, diag1, diag2)
                    # print (col,diag1,diag2)
                    bactrack(y+1, col_pos, n, col, diag1, diag2)
                    # backtrack to previous state
                    _remove(x, y, col_pos, col, diag1, diag2)
        
        ans = []
        col = [0] * n
        diag1 = [0] * (2*n-1)
        diag2 = [0] * (2*n-1)
        bactrack(0, [], n, col, diag1, diag2)
        return ans