Solution to Google Code Jam World Finals 2017 - Problem A. Dice Straight in Python

solution.py

import itertools
import numpy as np
from scipy.optimize import linear_sum_assignment

with open('A-small-practice.in', 'rt') as f:
    T = int(next(f))
    for t in range(T):
        dice = []
        N = int(next(f))
        for n in range(N):
            die = next(f)
            dice.append(list(map(int, die.split())))

        storage = {}
        for i, die in enumerate(dice):
            for face in die:
                if face in storage:
                    storage[face].append(i)
                else:
                    storage[face] = [i]

        longest_straight = 1
        faces = sorted(storage.keys())
        for i in range(len(faces)):
            for length in range(longest_straight+1, N+1):
                sequence = faces[i:i+length]
                n = len(sequence)
                if n < length:
                    break
                if sequence == list(range(sequence[0], sequence[-1]+1)):
                    dices = [storage[face] for face in sequence]
                    set_dices = list(set(itertools.chain(*dices)))

                    # Check if there are enough dices to make a sequence
                    if len(set_dices) < len(sequence):
                        break

                    # Bipartite matching
                    cost = np.zeros((len(set_dices),n))
                    for c_j, face in enumerate(sequence):
                        for dice in storage[face]:
                            cost[set_dices.index(dice), c_j] = -1
                    row_ind, col_ind = linear_sum_assignment(cost)

                    # Check if result is possible
                    if np.count_nonzero([set_dices[row_i] in dices[col_i] for row_i, col_i in zip(row_ind, col_ind)]) != n:
                        break

                    longest_straight = n
                else:
                    break

        print('Case #%d: %d' % (t+1, longest_straight))