AoC 2021 - Day 6

day06.ipynb

day6.py

input = [3, 4, 3, 1, 2]

def number_in_generation(generation, days):
    if days < 9 * generation:
        return 0

    steps = (days - 9 * generation) // 7 + 1
    count = 1
    for i in range(0, generation + 1):
        count = count * (steps + i) // (1 + i)

    return count

def calculate_total(days):
    count = 1
    for gen in range(days):
        num = number_in_generation(gen, days)
        if num == 0:
            break
        count += num
    return count

def calculate_with_offset(startvalue, days):
    days -= startvalue
    if days <= 0:
        return 1
    else:
        return calculate_total(days-1)

def calculate(input, days):
    counts = [calculate_with_offset(x, days) for x in input]
    return sum(counts)

part1 = calculate(input, 80)
part2 = calculate(input, 256)

print("Part 1:", part1)
print("Part 2:", part2)