connect_four.py

def coding_challenge_specific_io(challenge_function):
    with open('input') as in_data:
        with open('output', 'w') as out_data:
            in_lines = [x.strip() for x in in_data.readlines()[1:]]
            in_data = parse_input_data(in_lines)
            challenge_complete_lines = [challenge_function(data) for data in in_data]
            out_lines = ['Case #{}: {}\n'.format(i+1, line) for i, line in enumerate(challenge_complete_lines)]
            out_data.writelines(out_lines)


def parse_input_data(in_lines):
    in_blocks = []
    in_buffer = {'matches': 0, 'block': []}
    for in_line in in_lines:
        if in_line[0].isdigit():
            in_blocks.append(in_buffer)
            in_buffer = {'matches': int(in_line[-1]), 'block':[]}
        else:
            in_buffer['block'].append(in_line)
    in_blocks.append(in_buffer)
    return in_blocks[1:]


def shift_line_right(line):
    squished_line = line.replace('.', '')
    removed_dots_count = len(line) - len(squished_line)
    return '.' * removed_dots_count + squished_line


def shift_right(block):
    return [shift_line_right(x) for x in block]


def check_horizontal(block_data):
    m = block_data['matches']
    for line in block_data['block']:
        if 'B' * m in line or 'R' * m in line:
            return True
    return False


def rotate(block):
    return [''.join(i) for i in zip(*block[::-1])]


def check_vert(block_data):
    b = block_data.copy()
    b['block'] = rotate(b['block'])
    return check_horizontal(b)


def check_top_half_diag(block, matches):
    for y in range(len(block)):
        diag_line = ''.join([block[y-x][x] for x in range(y+1)])
        if 'B' * matches in diag_line or 'R' * matches in diag_line:
            return True
    return False


def check_diag(block_data):
    b = block_data['block']
    for _ in range(4):
        if check_top_half_diag(b, block_data['matches']):
            return True
        b = rotate(b)
    return False


def check_all(block_data):
    checkers = [check_horizontal, check_vert, check_diag]
    return any([x(block_data) for x in checkers])


def remove_color(letter, block_data):
    block_data = block_data.copy()
    block_data['block'] = [line.replace(letter, '.') for line in block_data['block']]
    return block_data


def check_for_blue_win(block_data):
    blue_block_data = remove_color('R', block_data)
    return check_all(blue_block_data)


def check_for_red_win(block_data):
    red_block_data = remove_color('B', block_data)
    return check_all(red_block_data)


def connect_four_solver(block_data):
    block_data['block'] = shift_right(block_data['block'])
    blue = check_for_blue_win(block_data)
    red = check_for_red_win(block_data)
    if all([red, blue]):
        return "both"
    elif red:
        return "red"
    elif blue:
        return "blue"
    else:
        return "neither"


if __name__ == '__main__':
    coding_challenge_specific_io(connect_four_solver)