advent_of_code_2021

day06_part1.py (3787B)

---

 #!/usr/bin/env python
 """Advent of Code 2021, day 6 (part 1): simulate an exponentially growing
 school of lanternfish using an initial list of fish and their ages"""
 
 # For this puzzle, the most important decision is the data structure to use to
 # represent the school of fish. Since there's exponential growth, any scheme
 # that uses an element for each fish could quickly get out of hand. Since the
 # set of possible ages is quite small, it's much more efficient to use counters
 # for each age. I'm using a pandas Series because of course I am (but also, I'm
 # confronting my aversion to pandas index attributes head on), but a simple
 # dict would work fine here too, or even a list if you're careful with updating
 # indices.
 
 import functools
 from utils import get_puzzle_input, convert_int_line_to_series
 
 def convert_input_to_series(input_string):
     """Parse the input (a sequence of fish with their age in days) and return a
     Series, where the indices represent fish ages (in days) and the values
     represent the count of fish that age"""
     return convert_int_line_to_series(input_string).value_counts()
 
 # For this implementation, we'll advance the fish in age and spawn new fish all
 # in one step--even though, conceptually, it would be nice to separate these
 # steps--because that way it'll be easier to separate the fish that are ready
 # to spawn from the babies who aren't.
 
 def step_one_day(fish_counts):
     """Given a Series of fish counts, simulate one day of ageing/growth of the
     school, and return a new Series"""
     # Once again, I'm trying to program more functionally, which means I'm
     # treating data structures as immutable (outside the function) even when it
     # might be more efficient to update things in-place. Theoretically we're
     # trading efficiency for "safety", although I'm perfectly aware that for
     # these puzzles it really doesn't matter. But inside the actual function, I
     # am editing things in-place, so that's a bit of a fail.
     # Implementation note: I looked at a couple methods of updating the axis
     # and was surprised by the results.
     # %timeit fish_counts.set_axis((fish_counts.index - 1) % 9)
     #> 305 µs ± 6.05 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
     # %timeit fish_counts.rename(lambda x: (x - 1) % 9)
     #> 93.1 µs ± 2.19 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
     updated_fish_counts = fish_counts.rename(lambda x: (x - 1) % 9)
     # We're cheating a bit here. Rather than resetting an adult's counter back
     # to 6 and then adding babies with a counter set to 8, we're resetting
     # adults' counters to 8 (representing the babies) and then adding the
     # adults back into the count. That's why we take the counter value modulo 9
     # instead of the obvious 7.
     if 8 in updated_fish_counts:
         if 6 in updated_fish_counts:
             updated_fish_counts[6] += updated_fish_counts[8]
         else:
             updated_fish_counts[6] = updated_fish_counts[8]
     return updated_fish_counts
 
 def count_fish_after(fish_counts, num_days):
     """Given a Series of fish counts and a number of days, return the total
     number of fish in the school as an integer"""
     return int(
         functools.reduce(lambda x, y: step_one_day(x),
                          range(num_days),
                          fish_counts)
         .sum()
     )
 
 def solve_puzzle(input_string):
     """Return the numeric solution to the puzzle"""
     return count_fish_after(convert_input_to_series(input_string), 80)
 
 def main():
     """Run when the file is called as a script"""
     assert solve_puzzle("3,4,3,1,2\n") == 5934
     print("Number of lanternfish after 80 days:",
           solve_puzzle(get_puzzle_input(6)))
 
 if __name__ == "__main__":
     main()