day6.py

from time import time
import numpy as np
from numpy.linalg import matrix_power
from collections import Counter

start = time()

nums = [0] * 9

for idx, count in Counter(map(int, open("priv/inputs/year2021/day6.txt").readlines()[0].split(","))).items():
    nums[idx] = count

nums = np.array(nums)

matrix = np.array([
    [0, 0, 0, 0, 0, 0, 1, 0, 1],
    [1, 0, 0, 0, 0, 0, 0, 0, 0],
    [0, 1, 0, 0, 0, 0, 0, 0, 0],
    [0, 0, 1, 0, 0, 0, 0, 0, 0],
    [0, 0, 0, 1, 0, 0, 0, 0, 0],
    [0, 0, 0, 0, 1, 0, 0, 0, 0],
    [0, 0, 0, 0, 0, 1, 0, 0, 0],
    [0, 0, 0, 0, 0, 0, 1, 0, 0],
    [0, 0, 0, 0, 0, 0, 0, 1, 0],
], dtype=object)

dot = matrix_power(matrix, 256).transpose().dot(nums)

print(dot.sum())

print(f"Took: {time() - start}s")