LC 2056 [H] Number of Valid Move Combinations On Chessboard - ALawliet/algorithms GitHub Wiki

>>> s = {1,2,3,4,5}
>>> s + {6,7}
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: unsupported operand type(s) for +: 'set' and 'set'
>>> s - {5}
{1, 2, 3, 4}
>>> s | {4,5}
{1, 2, 3, 4, 5}
>>> s & {4,5}
{4, 5}

LINE = [(0, 1), (1, 0), (0, -1), (-1, 0)]
DIAG = [(1, 1), (1, -1), (-1, 1), (-1, -1)]         

class Solution:
    def countCombinations(self, pieces: List[str], positions: List[List[int]]) -> int:
        board = [[set() for _ in range(8)] for _ in range(8)]
        n = len(pieces)
        
        for pos in positions:
            # adjust the positions to make them start with 0 
            pos[0] -= 1
            pos[1] -= 1
        
        all_time = set(range(1, 8)) # move for at most 8 seconds because only 8x8 board
        
        def branch(i):
            if i == n:
                # if the last piece has a valid destination, count this move combination
                return 1
            
            ans = 0
            r, c = positions[i]
            # print(board[r][c], all_time, board[r][c] & all_time)
            if not board[r][c] & all_time:
                # if the position of piece i is available all the time, consider not moving this piece at all
                # print(board[r][c], all_time, board[r][c] | all_time)
                board[r][c] |= all_time
                ans += branch(i + 1) # try to find the move of the next piece
                board[r][c].clear() # undo the move for the current piece
                
            DIRS = [] # the directions the current piece might move to
            
            if pieces[i] in ("queen", "rook"):
                DIRS += LINE
            if pieces[i] in ("queen", "bishop"):
                DIRS += DIAG              
            
            for dr, dc in DIRS:
                # consider a direction the current piece moves to
                # the current location of the move
                nr = r + dr
                nc = c + dc
                count = 1 #[#the](https://leetcode.com/problems/valid-parentheses) time counter
                while 0 <= nr < 8 and 0 <= nc < 8 and count not in board[nr][nc]:
                    # if board[x][y] is available at time count, the move is valid
                    board[nr][nc].add(count) 
                    count += 1
                    rest = set(range(count, 8))
                    # print(board[nr][nc], rest, board[nr][nc] & rest)
                    if not board[nr][nc] & rest:
                        # if the board is available for the rest of the time, the move might stop here
                        board[nr][nc] |= rest # in-place or, union
                        # print(board[nr][nc], rest, board[nr][nc] | rest)
                        ans += branch(i + 1)
                        # print(board[nr][nc], rest, board[nr][nc] - rest)
                        board[nr][nc] -= rest
                    nr += dr
                    nc += dc
                    
                # undo the last update
                count -= 1
                nr -= dr
                nc -= dc
                
                while count:
                    # undo the updates of the piece
                    board[nr][nc].remove(count)
                    count -= 1
                    nr -= dr
                    nc -= dc
                    
            return ans
        
        return branch(0)