The main things to deal with in this problem are:
We utilize a collection of dictionaries and functions to help us with these tasks are outlined below.
The board is stored as a list of length 16, where index 0 is referring to the upper left-hand square
of the simplified $4 \times 4$ board, and index 15 is the lower right-hand square.
Note that gathering the input is done here in a different fashion than is done on HackerRank.
import sys
# Create piece and board dictionaries.
# Used for translating the input of the problem into more
# usable values for the remaining functions.
PIECE_DICT = {'Q' : 4, 'R' : 3, 'B' : 2, 'N' : 1}
BOARD_DICT = {'A': {'1' : 0, '2' : 1, '3' : 2, '4' : 3},
'B': {'1' : 4, '2' : 5, '3' : 6, '4' : 7},
'C': {'1' : 8, '2' : 9, '3' : 10, '4' : 11},
'D': {'1' : 12, '2' : 13, '3' : 14, '4' : 15},
}
MOVE_LIST = ['r','d','l','u','dr','dl','ur','ul']
MOVE_DICT = dict((x,[-1]*16) for x in MOVE_LIST)
def get_init_board(n_white, n_black):
"""
Initialize board
"""
board = [0] * 16
for i in range(n_white + n_black):
pc, row, col = sys.stdin.readline().split()
piece = PIECE_DICT[pc] * (-1)**(0+(i>=n_white))
position = BOARD_DICT[row][col]
board[position] = piece
return board
def precompute_moves():
"""
Determine valid move directions for each position on the board
"""
for i in range(16):
okright = (i % 4 != 3)
okleft = (i % 4 != 0)
okup = (i >= 4)
okdown = (i <= 11)
# Rightward moves
if okright:
MOVE_DICT['r'][i] = i + 1
# Leftward moves
if okleft:
MOVE_DICT['l'][i] = i - 1
# Downward moves
if okdown:
MOVE_DICT['d'][i] = i + 4
# Upward moves
if okup:
MOVE_DICT['u'][i] = i - 4
# Right-down diagonal moves
if okright and okdown:
MOVE_DICT['dr'][i] = i + 5
# Left-down diagonal moves
if okleft and okdown:
MOVE_DICT['dl'][i] = i + 3
# Right-up diagonal moves
if okright and okup:
MOVE_DICT['ur'][i] = i - 3
# Left-up diagonal moves
if okleft and okup:
MOVE_DICT['ul'][i] = i - 5
def n_moves(board,i,rv):
"""
Determine valid knight moves
"""
for s in range(4):
j = MOVE_DICT[MOVE_LIST[s]][i]
if j == -1:
continue
j = MOVE_DICT[MOVE_LIST[s]][j]
if j == -1:
continue
for t in range(2):
tt = (s + 2 * t + 1) % 4
k = MOVE_DICT[MOVE_LIST[tt]][j]
if k == -1:
continue
if board[k] <= 0:
x = board[:]
x[i] = 0
x[k] = 1
rv.append(x)
def dir_moves(board,i,direction,rv):
"""
Determine valid single-directional moves (for all pieces ex. knights)
"""
pos = i
while True:
pos = MOVE_DICT[direction][pos]
if pos == -1:
return
if board[pos] > 0:
return
x = board[:]
x[pos] = x[i]
x[i] = 0
rv.append(x)
if board[pos] < 0:
return
def b_moves(board,i,rv):
# Bishops only move diagonally (move categories 4-7)
for d in MOVE_LIST[4:]:
dir_moves(board,i,d,rv)
def r_moves(board,i,rv):
# Rooks only move horizontally (move categories 0-3)
for d in MOVE_LIST[0:4]:
dir_moves(board,i,d,rv)
def q_moves(board,i,rv):
for d in MOVE_LIST:
dir_moves(board,i,d,rv)
def new_boards(board):
rv = []
for i in range(16):
if board[i] == 1:
n_moves(board,i,rv)
if board[i] == 2:
b_moves(board,i,rv)
if board[i] == 3:
r_moves(board,i,rv)
if board[i] == 4:
q_moves(board,i,rv)
return rv
def is_win(board):
for p in board:
if p == -4:
return False
return True
memodict = {}
def board_value(board, depth):
state = tuple(board + [depth])
if state in memodict:
return memodict[state]
if depth == 0:
return 0
bb = new_boards(board)
if len(bb) == 0:
return 0
for x in bb:
if is_win(x):
memodict[state] = 1
return 1
best = -1
for x in bb:
y = [-p for p in x]
v = -board_value(y, depth - 1)
if v > best:
best = v
if best == 1:
break
memodict[state] = best
return best
if __name__ == '__main__':
for _ in range(int(sys.stdin.readline())):
n_white, n_black, depth = (int(x) for x in sys.stdin.readline().split())
board = get_init_board(n_white, n_black)
precompute_moves()
value = board_value(board, depth)
if value == 1:
print("YES")
else:
print("NO")