Learning Python with Advent of Code Walkthroughs
Dazbo's Advent of Code solutions, written in Python
Advent of Code 2021 - Day 6
Useful Links
Concepts and Packages Demonstrated
Problem Intro
Here I’ve documented three different approaches to this problem:
- Solution #1 - The naive way
- Solution #2 - The NumPy way
- Solution #3 - Using
deque.rotate()
We’re told that lanternfish populations grow exponentially, and we’re asked to model the growth of the population. Each lanternfish spawns a new lanternfish every 7 days, except for new lanternfish, which require an extra 2 days before spawning a new fish.
Our input data is a list of numbers, which are the number of days left until each fish spawns another fish. The data looks like this:
3,4,3,1,2
Part 1
Given our input data, how many lanternfish will there be after 80 days?
Solution #1
At first glance, this seems like a simple problem. The naive approach is short and easy to write. Alas, as we’ll see, it has a problem!
We create list of fish_timers from the input data, where each timer is the number of days until the next spawn. With each day, we update the existing fish timers, and add the new timers.
input_file = os.path.join(SCRIPT_DIR, INPUT_FILE)
with open(input_file, mode="rt") as f:
data = f.read()
fish_timers = list(map(int, data.split(",")))
logger.debug(fish_timers) # E.g. [3,4,3,1,2]
fish_timers_copy = fish_timers.copy()
days = 80
for _ in range(1, days+1):
# we can't enumerate since we don't extra iterations when we spawn
for i in range(len(fish_timers_copy)):
if fish_timers_copy[i] == 0:
fish_timers_copy[i] = 6 # reset this fish timer
fish_timers_copy.append(8) # spawn a new fish
else:
fish_timers_copy[i] -= 1 # decrement this fish timer
logger.info("After %d days, there are %d fish", days, len(fish_timers_copy))
The code above reads in the input data, splits it at the commas, and then converts each value to an int. Nothing new there.
First, we make a copy of the fish_timers. This is just so that we can use the same original input list when we do part 2.
Then we iterate through 80 days. Note that in this for loop, we store each day in the variable _. This is a convention in Python: if we’re not interested in this variable, and we’re not going to use it, then we use the value _.
Then, for each day, we now iterate through each fish_timer (i.e. each fish), and do the following:
- If the fish timer is at 0, the fish must spawn a new fish.
- Add a new timer with a value of 8. (Since all new fish require 9 days before spawning a new fish. Timers are 0-indexed.)
- Change this fish timer to 6, i.e. since this fish is now 7 days away from its next spawn.
- For all other timers, decrement their value by 1.
And for Part 1, this approach works fine!
2022-01-09 22:11:00.555:INFO:__main__: After 80 days, there are 376194 fish
Part 2.
Now we’re simply asked to determine how many fish there will be after 256 days.
This presents us with a predictable problem. Our naive solution doesn’t scale!
| Number of days | Time to solve |
|---|---|
| 80s | 0.5s |
| 90s | 1.1s |
| 100s | 3.1s |
| 110s | 6.7s |
| 120s | 17s |
If we extrapolate, this is going to take days to run! We need a different approach.
Solution #2
At this point, it’s obvious we can’t model every fish. We actually only need to track how many there of each timer, e.g. how many fish are spawning now, how many are 1 day away, how many are 2 days away, etc.
Solving the Problem
Here we can read the source data directly into a numpy array, using np.loadtxt. We split the numbers at the commas, and convert each number to an int using dtype=np.int8.
input_file = os.path.join(SCRIPT_DIR, INPUT_FILE)
# read input, which is csv of ints
# Each value is a fish timer, i.e. days until spawning another fish
data = np.loadtxt(input_file, delimiter=",", dtype=np.int8) # [3 4 3 1 2]
Then we use np.unique() with the return_counts=True attribute to return all the unique fish timers in the input data, and count them all.
# np.unique() returns ordered unique items. With return_counts=True, it includes counts.
# Timers [1 2 3 4], counts [1 1 2 1]
init_fish_timers, counts = np.unique(data, return_counts=True)
We know there will be timers for all values 0-8 inclusive. Our input data may not include all of these types of timers. So let’s create a one dimensional numpy array of size 9 (i.e. 0-8 inclusive), and initialise all the values to 0. This represents our counters for each fish timer.
Then we can update this ‘zero’ array using the counts we previously obtained from the input data. This works by using the init_fish_timers to find all the corresponding index positions in fish_timers, and then updating value at that position by using the corresponding position in the counts array.
# Initialise fish timers, by setting index positions to the counts in the array
fish_timers = np.zeros(9, dtype=np.uint64) # [0 0 0 0 0 0 0 0 0]
fish_timers[init_fish_timers] = counts # [0 1 1 2 1 0 0 0 0]
Now we create a get_fish_count() function, which returns the number of fish after the required number of days. It works by:
- Looping through the required number of days.
- Calling
np.roll(fish, -1), which effectively shifts the entire array one to the left, and wraps items on the far left over to the right. Thus, timers that were 0 now become 8 (representing new fish), items that were 1 become 0, etc. - We increment the number of timers at
[6]by the number of fish that are now at[8]. I.e. the number of fish that have reset is equal to the number of fish that were just spawned.
def get_fish_count(fish_timers: np.ndarray, day_num: int) -> int:
fish = np.copy(fish_timers) # create a new copy so we don't mutate the original fish
for _ in range(day_num):
fish = np.roll(fish, -1) # Roll: 2 becomes 1, 1 becomes 0, 0 becomes 8 (spawned fish)
fish[6] += fish[8] # Add fish that were 0 are now reset to 6.
return sum(fish)
We can now call this function with both 80 and 256, thus solving both Part 1 and Part 2:
for day in (80, 256):
logger.info("At day %d, count=%d", day, get_fish_count(fish_timers, day))
The output looks like this:
2022-01-09 22:58:19.820:INFO:__main__: At day 80, count=376194
2022-01-09 22:58:19.830:INFO:__main__: At day 256, count=1693022481538
2022-01-09 22:58:19.831:INFO:__main__: Execution time: 0.0353 seconds
It’s quite quick!
Solution #3
This solution is super fast!
We don’t need numpy. We can achieve the same results using a Counter and a deque.
First, read the data, split it at the commas, convert each value to int, and then store the counts of each type of fish timer:
input_file = os.path.join(SCRIPT_DIR, INPUT_FILE)
with open(input_file, mode="rt") as f:
data = f.read()
data = [int(x) for x in data.split(",")]
fish_timer_counts = Counter(data) # count the timers for each fish
Now we create a deque that contains our nine timer values, 0-8 inclusive. The values are initialised to 0 if there weren’t any in the input data. Else, we use the counts from the input data.
# initialise an array of timer counts, for all timer values 0-8
timers = deque()
for timer in range(9):
timers.append(fish_timer_counts[timer] if fish_timer_counts[timer] else 0)
As before, we have a get_fish_count() function. It does the following:
- Iterates through the required number of days.
- Rotates the entire deque to the left, i.e. so that 2 becomes 1, 1 becomes 0, 0 becomes 8, etc.
- Increments the count of [6] by the number of fish at [8], i.e. the count of fish that just reset.
for days in (80, 256):
logger.info("Fish at day %d: %d", days, get_fish_count(timers, days))
def get_fish_count(timers: deque, days: int):
fish = timers.copy() # just so we can repeat this method with a different # of days
for _ in range(days):
fish.rotate(-1)
fish[6] += fish[8] # count of newly spawned fish is same as count of fish that need to be reset
return sum(fish)
Output:
2022-01-09 22:59:06.606:INFO:__main__: Fish at day 80: 376194
2022-01-09 22:59:06.607:INFO:__main__: Fish at day 256: 1693022481538
2022-01-09 22:59:06.608:INFO:__main__: Execution time: 0.0027 seconds
Yep, under 3ms. Pretty quick!