[LeetCode] 51. N-Queens

Problem : https://leetcode.com/problems/n-queens/

A typical backtracking problem. Simulate the process to check if it is valid to put queen on current position. Otherwise, return to the last state.

 

class Solution:
    def solveNQueens(self, n: int) -> List[List[str]]:
        board = [[0] * n for _ in range(n)]
        result = []
        
        def toBoard(queens):
            return [''.join(['Q' if x == queens[y] else '.' for x in range(n)]) for y in range(n)]
        
        @cache
        def getOccupied(y, x):
            occupied = [(y,x)]
            
            # because we scan from left to right and top to bottom
            # it only needs to mark the poisition on right and on rows below
            for dy, dx in [(1,0), (0,1), (1, 1), (1, -1)]:
                tmpY, tmpX = y + dy, x + dx
                while 0 <= tmpY < n and 0 <= tmpX < n:
                    occupied.append((tmpY, tmpX))
                    tmpY += dy
                    tmpX += dx
                
            return occupied
        
        def backtrack(y, queens):
            if y == n:
                yield toBoard(queens)
            else:
                for x in range(n):
                    if board[y][x]:
                        continue
                    
                    myOccupied = getOccupied(y, x)
                    
                    for occupiedY, occupiedX in myOccupied:
                        board[occupiedY][occupiedX] += 1
                        
                    queens.append(x)
                    yield from backtrack(y+1, queens)
                    queens.pop()
                    
                    for occupiedY, occupiedX in myOccupied:
                        board[occupiedY][occupiedX] -= 1
        
        return list(backtrack(0, []))

Edited on 05/22/2021. Simplify the method for caculating occupied positions.

Edited on 05/29/2021. Optimize performance by only marking poisition on right and on rows below.