advent_of_code/6/6.py

#!/usr/bin/env python
# 2021 - Advent Of Code - 6

def parse_file(file):
    with open(file) as f:
        line = f.readline()
    return [int(x) for x in line.split(',')]


def iterate(state):
    new_fish = 0
    for i, _ in enumerate(state):
        if state[i] == 0:
            new_fish += 1
            state[i] = 6
        else:
            state[i] -= 1
    state.extend([8]*new_fish)


init_state = parse_file('input.txt')
print(f'Initial state: {init_state}')

DAYS = 80

for d in range(DAYS):
    iterate(init_state)
    # print(f'After {d+1:2} days: {init_state}')

print(f'There\'s {len(init_state)} lanterfish after {DAYS} days')