[LeetCode] 52. N-Queens II

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

Same to the 51. N-Queens

class Solution:
    def totalNQueens(self, n: int) -> int:
        
        @cache
        def getOccupies(y, x):
            occupies = [(y, x)]
            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:
                    occupies.append((tmpY, tmpX))
                    tmpY += dy
                    tmpX += dx
            
            return occupies
        
        
        def backtrack(y, board):
            if y == n:
                yield 1
            else:
                for x in range(n):
                    if board[y][x]:
                        # skip this position as it has been occupied
                        continue
                    
                    myOccupies = getOccupies(y,x)
                    
                    for py, px in myOccupies:
                        board[py][px] += 1
                    
                    yield from backtrack(y+1, board)
                    
                    for py, px in myOccupies:
                        board[py][px] -= 1
        
        board = [[0] * n for _ in range(n)]
        return sum(backtrack(0, board))

Edited on 05/22/2021. Use memorization to simplify the code for checking whether one position is valid to put a Queen.

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