Advent of Code 2021 Day 6 in Python

day06.py

from collections import deque
import itertools

def init(fish):
    """
    >>> init([3, 4, 3, 1, 2])
    [0, 1, 1, 2, 1, 0, 0, 0, 0]
    """
    result = [0] * 9
    for day, fish in itertools.groupby(sorted(fish)):
        result[day] = len(list(fish))
    return result

def step(state):
    """
    >>> step(deque([0, 1, 1, 2, 1, 0, 0, 0, 0]))
    deque([1, 1, 2, 1, 0, 0, 0, 0, 0])
    >>> step(deque([1, 1, 2, 1, 0, 0, 0, 0, 0]))
    deque([1, 2, 1, 0, 0, 0, 1, 0, 1])
    """
    spawning = state.popleft()
    state.append(spawning) # new fish
    state[6] += spawning # old fish
    return state

def multistep(state, days):
    """
    >>> sum(multistep(deque([0, 1, 1, 2, 1, 0, 0, 0, 0]), 18))
    26
    >>> sum(multistep(deque([0, 1, 1, 2, 1, 0, 0, 0, 0]), 80))
    5934
    >>> sum(multistep(deque([0, 1, 1, 2, 1, 0, 0, 0, 0]), 256))
    26984457539
    """
    for day in range(days):
        step(state)
    return state    

if __name__ == "__main__":
    from doctest import testmod
    testmod()

    state = init([int(i) for i in open("day06.txt").read().split(",")])
    print(sum(multistep(deque(state), 256)))