advent_of_code_exercice_6.py

import sys
from numpy import *
from collections import Counter

n = int(sys.argv[1])
def mat_fun(i,j): return (j == (i+1)%9) | ((i==6) & (j==0))
M = matrix(fromfunction(mat_fun, (9, 9)), dtype=int)**n
ages=Counter(map(int,input().split(',')))
sizes=array([ages[i] for i in range(9)])
print(M.dot(sizes).sum())