yeti mini chess | Python Fiddle

# Enter your code here. Read input from STDIN. Print output to STDOUT
g = int(raw_input())
blarg = 0
translate = {'1':0,'A':0,'2':1,'B':1,'3':2,'C':2,'4':3,'D':3}
seen = {}
def get_legal_positions(pieces, piece): 
    p = [pc for pc in list(pieces) if pc[2] != piece[2]]
    res = []
    new_squares = []
    
    square = [piece[0],piece[1]]
    i,j = square[0],square[1]
    #Queen
    if 'Q' in piece[2]:
        new_squares = new_squares + [(i+x,j) for x in range(1,4) if i+x<4 ] + [(i-x,j) for x in range(1,4) if i-x>=0]
        new_squares = new_squares + [(i,j+x) for x in range(1,4) if j+x<4 ] + [(i,j-x) for x in range(1,4) if j-x>=0]
        new_squares = new_squares + [(i+x,j+x) for x in range(1,4) if i+x<=j+x<4 ] + [(i-x,j-x) for x in range(1,4) if i-x<=j-x>=0]
        new_squares = new_squares + [(i-x,j+x) for x in range(1,4) if i-x>=0 and j+x<4 ] + [(i+x,j-x) for x in range(1,4) if i+x<4 and j-x>=0]
    #Rook
    elif 'R' in piece[2]:
        new_squares = new_squares + [(i+x,j) for x in range(1,4) if i+x<4 ] + [(i-x,j) for x in range(1,4) if i-x>=0]
        new_squares = new_squares + [(i,j+x) for x in range(1,4) if j+x<4 ] + [(i,j-x) for x in range(1,4) if j-x>=0]
    #Bishop
    elif 'B' in piece[2]:
        new_squares = new_squares + [(i+x,j+x) for x in range(1,4) if i+x<=j+x<4 ] + [(i-x,j-x) for x in range(1,4) if i-x<=j-x>=0]
        new_squares = new_squares + [(i-x,j+x) for x in range(1,4) if i-x>=0 and j+x<4 ] + [(i+x,j-x) for x in range(1,4) if i+x<4 and j-x>=0]
    #Knight
    elif 'N' in piece[2]:
        if i-2 >= 0:
            if j+1<4:
                new_squares.append((i-2,j+1))
            if j-1>=0:
                new_squares.append((i-2,j-1))
        if i+2 >= 0:
            if j+1<4:
                new_squares.append((i+2,j+1))
            if j-1>=0:
                new_squares.append((i+2,j-1))
        if j-2 >= 0:
            if i+1<4:
                new_squares.append((i+1,j-2))
            if i-1>=0:
                new_squares.append((i-1,j-2))
        if j+2 >= 0:
            if i+1<4:
                new_squares.append((i+1,j+2))
            if i-1>=0:
                new_squares.append((i-1,j+2))
    for s in new_squares:
        new_pieces = [pc for pc in p]
        if any(map(lambda x: x[0] == s[0] and x[1] == s[1], p)):
            new_new_pieces = []
            for n in new_pieces:
                if n[0] == s[0] and n[1] == s[1]:
                    new_new_pieces.append((n[0],n[1],piece[2]))
                else:
                    new_new_pieces.append(n)
            new_pieces = [n for n in new_new_pieces]
        else:
            new_pieces.append((s[0],s[1],piece[2]))
        new_pieces.sort()
        res.append(tuple(new_pieces))
    return res
def mate_in_x(pieces, whose_turn, moves_left):
    if (pieces, whose_turn) in seen:
        result_of_position = seen[(pieces, whose_turn)]
        if moves_left >= result_of_position[0] and result_of_position[1]:
            res = True    
        else:
            res = False
    elif all(map(lambda x: x[2] != 'QW', pieces)):
        res = False
    elif all(map(lambda x: x[2] != 'QB', pieces)):
        res = True
    elif moves_left == 0:
        res = all(map(lambda x: x[2] != 'QB', pieces))
    else:
        if whose_turn == 'white':
            res = False
            wpieces = [wp for wp in pieces if 'W' == wp[2][1]]
            for wp in wpieces:
                for p in get_legal_positions(pieces,wp):                    
                    res = res or mate_in_x(p,'black',moves_left-1)
                    if res:
                        if (p, 'white') not in seen:
                            seen[(p, 'white')] = [moves_left, res]
                        else:
                            result_of_position = seen[(p, 'white')]
                            if result_of_position[0] < moves_left:
                                result_of_position = [moves_left, True]
                            elif not result_of_position[1]:
                                result_of_position[1] = True
                        return res
        else:
            res = True
            bpieces = [bp for bp in pieces if 'B' == bp[2][1]]
            for bp in bpieces:
                for p in get_legal_positions(pieces,bp):
                    res = res and mate_in_x(p,'white',moves_left)
                    if not res:
                        if (p, 'black') not in seen:
                            seen[(p, 'black')] = [moves_left, res]
                        return res
    if (pieces, whose_turn) not in seen:
        seen[(pieces, whose_turn)] = [moves_left, res]
    return res
for game in range(g):    
    pieces = []
    w,b,m = map(int,raw_input().strip().split())
    for W in range(w):
        p,i,j = raw_input().strip().split()
        pieces.append((translate[i],translate[j],p + 'W')) 
    for B in range(b):
        p,i,j = raw_input().strip().split()
        pieces.append((translate[i],translate[j],p + 'B'))
    pieces = tuple(pieces)
    if mate_in_x(pieces, 'white', 0, m) <= m:
        print 'YES'
    else:
        print 'NO'

Python Fiddle

Python Cloud IDE