advent-of-code/2021/day6.py

import sys

values = [int(c) for c in sys.stdin.read().strip().split(",")]

days = 256
lanterns = {day: 0 for day in range(days)}
for value in values:
    for day in range(value, days, 7):
        lanterns[day] += 1

for day in range(days):
    for day2 in range(day + 9, days, 7):
        lanterns[day2] += lanterns[day]

# part 1
answer_1 = sum(v for k, v in lanterns.items() if k < 80) + len(values)
print(f"answer 1 is {answer_1}")

# part 2
answer_2 = sum(lanterns.values()) + len(values)
print(f"answer 2 is {answer_2}")