advent_of_code_2021

My attempts to work through the 2021 Advent of Code problems.
git clone https://git.eamoncaddigan.net/advent_of_code_2021.git
Log | Files | Refs | README | LICENSE

day14_part1.py (4816B)

      1 #!/usr/bin/env python
      2 """Advent of Code 2021, day 14 (part 1): Extended Polymerization
      3 Perform 'pair insertions' on a 'polymer' string"""
      4 
      5 # Since the length of the string basically doubles every step of the problem,
      6 # we need to make sure we have an efficient representation of the data. And
      7 # once again, I'm going to use pandas (a Series) even when a dict would do just
      8 # fine so that I can continue to familiarize myself with it
      9 
     10 from typing import List, Tuple, cast
     11 from functools import reduce
     12 import pandas as pd
     13 from utils import get_puzzle_input
     14 
     15 EXAMPLE_INPUT = \
     16 """NNCB
     17 
     18 CH -> B
     19 HH -> N
     20 CB -> H
     21 NH -> C
     22 HB -> C
     23 HC -> B
     24 HN -> C
     25 NN -> C
     26 BH -> H
     27 NC -> B
     28 NB -> B
     29 BN -> B
     30 BB -> N
     31 BC -> B
     32 CC -> N
     33 CN -> C
     34 """
     35 
     36 def split_instructions(instructions_list: List[str]) -> pd.Series:
     37     """parses a list of the form ["AB", "C"] and returns a Series equivalent to
     38     [("A", "C"), ("C", "B")]"""
     39     return pd.Series([(instructions_list[0][0], instructions_list[1]),
     40                       (instructions_list[1], instructions_list[0][1])])
     41 
     42 def parse_input(input_string: str) -> Tuple[pd.Series, pd.Series, Tuple[str, str]]:
     43     """Read the polymer template and the pair insertion instructions and
     44     return:
     45     * A pandas series giving the count of each pair (pairs are the index
     46       elements and counts are the values)
     47     * A pandas series representing insertion instructions (start pair is
     48       the index and the values are new pairs)
     49     * The identity of the first and last elements in the chain, since every
     50       element but these will be double-counted"""
     51     # Again, python's built-in types are perfectly sufficient for this job, but
     52     # I'm still getting a handle on pandas's indexing.
     53     lines = input_string.rstrip('\n').split('\n')
     54 
     55     first_last = (lines[0][0], lines[0][-1])
     56 
     57     # Using a Series as a Tuple[str, str] -> int dictionary
     58     pair_counts = pd.Series(1,
     59                             index=zip(lines[0][0:-1], lines[0][1:]),
     60                             name='pair_count')
     61 
     62     # Using a Series to map Tuple[str, str] -> two Tuple[str, str] entries,
     63     # representing the new set of pairs created by an insertion
     64     instruction_pairs = pd.Series(lines[2:]).str.split(' -> ')
     65     pair_insertions = (
     66         instruction_pairs.apply(split_instructions)
     67         .set_axis(instruction_pairs.map(lambda x: (x[0][0], x[0][1])))
     68         .melt(value_vars=[0, 1], ignore_index=False, value_name='new_pairs')
     69         .loc[:, 'new_pairs']
     70     )
     71 
     72     return (pair_counts, pair_insertions, first_last)
     73 
     74 def insert_at_pairs(pair_counts: pd.Series,
     75                     pair_insertions: pd.Series) -> pd.Series:
     76     """Perform the pair insertions and return an updated series of pair
     77     counts"""
     78     # Merge the counts and insertions series to get a data frame
     79     return (
     80         pd.merge(pair_insertions, pair_counts,
     81                  left_index=True, right_index=True)
     82         .groupby('new_pairs')
     83         .sum()
     84         .rename_axis(None)
     85         .loc[:, 'pair_count']
     86     )
     87 
     88 def count_elements(pair_counts: pd.Series,
     89                    first_last: Tuple[str, str]) -> pd.Series:
     90     """Given the count of the pairs and the identity of the first and last
     91     element in the string, return the count of each individual element"""
     92     # Everything was double-counted except we're missing one instance each of
     93     # the first and last element (that's why we preserve them)
     94     # Note (to self?): `cast` is just here to make mypy happy
     95     element_count = (
     96         pd.concat([pair_counts.rename(lambda x: cast(tuple, x)[0]),
     97                    pair_counts.rename(lambda x: cast(tuple, x)[1])])
     98         .groupby(level=0)           # Group by indices
     99         .sum()
    100         .rename('element_count')    # Superflous, but nice
    101     )
    102     element_count[list(first_last)] += 1
    103     element_count /= 2
    104     return element_count
    105 
    106 def apply_insertions(pair_counts: pd.Series,
    107                      pair_insertions: pd.Series,
    108                      first_last: Tuple[str, str],
    109                      num_insertions: int) -> pd.Series:
    110     """Apply `num_insertions` insertions and get the final count of the
    111     _elements_"""
    112     return count_elements(
    113         reduce(lambda x, y: insert_at_pairs(x, pair_insertions),
    114                range(num_insertions),
    115                pair_counts),
    116         first_last
    117     )
    118 
    119 def solve_puzzle(input_string: str) -> int:
    120     """Return the numeric solution to the puzzle"""
    121     element_count = apply_insertions(*(parse_input(input_string) + (10,)))
    122     return int(element_count.max() - element_count.min())
    123 
    124 def main() -> None:
    125     """Run when the file is called as a script"""
    126     assert solve_puzzle(EXAMPLE_INPUT) == 1588
    127     print("Difference between most and least common elements:",
    128           solve_puzzle(get_puzzle_input(14)))
    129 
    130 if __name__ == "__main__":
    131     main()