Advent of Code 2021: Day 06 with Python, and numpy.roll()

Link to challenge on Advent of Code 2021 website

Loading the data

The data is a comma-separated list, so we can load this using numpy's np.loadtext() again. We set dtype to uint32 as we'll be using these numbers for indexes later and they'll need to be whole numbers.

lines = np.loadtxt("day_6.txt", delimiter=",", dtype="uint32")

Part 1

Part 1 asks us to track the number of exponentially reproducing fish. While it's possible at this stage to keep an array of every fish, since every fish of the same age behaves the same, we can just keep track of groups of them.

The demo data gives these fish ages: 3,4,3,1,2, this can be reduced to an array of fish, where the index is the age, and the value are their count: [0, 1, 1, 2, 1, 0, 0, 0]

The benefit of this approach is that when fish age, we can simply roll the array around:

Initial: [0, 1, 1, 2, 1, 0, 0, 0]
  Day 1: [1, 1, 2, 1, 0, 0, 0, 0]
  Day 2: [1, 2, 1, 0, 0, 0, 1, 1]
  Day 3: [2, 1, 0, 0, 0, 1, 2, 1]
                            ^  ^
  old fish get added here --^  ^
     new fish get added here --^

To assemble the fish, we can use numpy's np.unique() to return a frequency count of the fish:

fish = np.zeros(9)
age, count = np.unique(lines, return_counts=True)
fish[age] = count

And to implement the roll, we can use np.roll():

gen = np.copy(fish)
for _ in range(80):
    gen[7] += gen[0]
    gen = np.roll(gen, -1)
result = sum(gen)

Taking the sum of the array gives us the total fish.

Part 2

Exactly the same as Part 1, but simulated to 256 generations. I assume the reason for part 2 like this is to trip up anyone who implemented Part 1 inefficiently, since it would take a lot of RAM to keep track of individual fish without grouping them.

gen = np.copy(fish)
for _ in range(256):
    gen[7] += gen[0]
    gen = np.roll(gen, -1)
result = sum(gen)

Full Code

import numpy as np

lines = np.loadtxt("day_6.txt", delimiter=",", dtype="uint32")
fish = np.zeros(9)
age, count = np.unique(lines, return_counts=True)
fish[age] = count

gen = np.copy(fish)
for _ in range(80):
    gen[7] += gen[0]
    gen = np.roll(gen, -1)
result = sum(gen)
print("Part 1 result:", result)

gen = np.copy(fish)
for _ in range(256):
    gen[7] += gen[0]
    gen = np.roll(gen, -1)
result = sum(gen)
print("Part 2 result:", result)