Advent of Code using `numpy` & `pandas`

day01-part1.py

from pandas import read_csv

s = read_csv('data.txt', squeeze=True, header=None)
print(
    (s.diff() > 0).sum(),
)

day01-part2.py

from pandas import read_csv

s = read_csv('data.txt', squeeze=True, header=None)
print(
    (s.rolling(3).sum().diff() > 0).sum(),
)

day02-part1.py

from pandas import read_csv, Series, CategoricalDtype

d = CategoricalDtype('forward up down'.split())
m = Series({
    'forward': 1+0j,
    'up':      0-1j,
    'down':    0+1j,
}).pipe(lambda s: s.set_axis(s.index.astype(d)))

s = read_csv(
    'data.txt',
    sep=' ',
    names=['direction', 'distance'],
    index_col=['direction'], header=None,
    dtype={'direction': d},
    squeeze=True,
)

print(
    (pos := (s * m.loc[s.index].values).sum()).real * pos.imag,
)

day02-part2.py

from pandas import read_csv, Series, CategoricalDtype, DataFrame, IndexSlice

d = CategoricalDtype('forward up down'.split())
a = Series({'up': -1, 'down': 1}).pipe(lambda s: s.set_axis(s.index.astype(d)))
df = read_csv(
    'data.txt',
    sep=' ',
    names=['direction', 'distance'],
    header=None, dtype={'direction': d},
)
mask = df['direction'] == 'forward'
df = DataFrame({
    **{('horizontal', col): df[ mask][col] for col in df.columns},
    **{('vertical',   col): df[~mask][col] for col in df.columns},
}, index=df.index)

print(
    df.assign(
        aim=lambda df:
            df['vertical'].dropna()
            .pipe(lambda df: df['distance'] * a.loc[df['direction']].values)
            .cumsum()
            .reindex(df.index).pad().fillna(0)
    ).pipe(lambda df: df.join(
        DataFrame({
            ('vertical', 'position'):   df['horizontal', 'distance'] * df['aim'],
            ('horizontal', 'position'): df['horizontal', 'distance'],
        }, index=df.index).fillna(0).cumsum()
    )).iloc[-1].loc[IndexSlice[:, 'position', :]].product(),
)

day03-part1.py

from numpy import memmap, unique, int8, logspace
from pandas import DataFrame, Series

bits = [f'bit{idx}' for idx in range(12)]
dtype = [*{**{bit: int8 for bit in bits}, '_nl': int8}.items()]
xs = (memmap('data.txt', mode='r') == ord(b'1')).view(dtype)
axis = 'gamma epsilon'.split()

print(
    DataFrame({
        bit: Series(*unique(xs[bit], return_counts=True)).sort_index().set_axis(axis)
        for bit in bits
    })
    .pipe(lambda df: df * logspace(11, 0, 12, base=2, dtype=int))
    .sum(axis='columns').product()
)

day03-part2.py

from numpy import memmap, int8, logspace
from pandas import Series, DataFrame
from itertools import dropwhile

bits = [f'bit{idx}' for idx in range(12)]
dtype = [*{**{bit: int8 for bit in bits}, '_nl': int8}.items()]
xs = (memmap('data.txt', mode='r') == ord('1')).view(dtype)

def search(df, cols, pred):
    for col in cols:
        df = df[df[col].pipe(lambda bit:
            bit == bit.value_counts().reindex([0, 1], fill_value=0).pipe(pred)
        )]
        if not df.index.size:
            break
        yield df.index[0]

df = DataFrame(xs)
ratings = Series({
    'O₂': next(dropwhile(
        lambda df: df.size > 1,
        search(df, bits, lambda vc: vc[0] <= vc[1])
    )),
    'CO₂': next(dropwhile(
        lambda df: df.size > 1,
        search(df, bits, lambda vc: vc[0] > vc[1])
    )),
})

print(
    df.loc[ratings, bits]
      .pipe(lambda df: df * logspace(11, 0, 12, base=2, dtype=int))
      .sum(axis='columns').product()
)

day04.py

from itertools import groupby
from numpy import array, full, broadcast_to, swapaxes, invert
from xarray import DataArray

with open('data.txt') as f:
    draws = next(f).split(',')
    boards = groupby(f, lambda ln: not ln.strip())
    boards = (g for k, g in boards if not k)
    boards = [array([[int(x) for x in ln.split()] for ln in b]) for b in boards]
    boards = array(boards)
    draws = array([full(boards.shape[1:], int(x)) for x in draws])

draws = DataArray(
    data=draws,
    dims='round y x'.split(),
)
boards = DataArray(
    data=swapaxes(broadcast_to(boards, (len(draws), *boards.shape)), 0, 1),
    dims=['player', *draws.dims],
)

# play the game…
state = (boards == draws).cumsum(dim='round')
wins = (state.sum(dim='y') >= state.sizes['y']).sum(dim='x') + \
       (state.sum(dim='x') >= state.sizes['x']).sum(dim='y') > 0

all_winners = wins.argmax(dim='round')
winning_players = [all_winners.argmin(), all_winners.argmax()]
winning_rounds = all_winners[winning_players]

print(
    (
        (state.sel(player=winning_players, round=winning_rounds) == 0)
      * boards.sel(player=winning_players, round=winning_rounds)
    ).sum(dim=['x', 'y']) * draws.sel(round=winning_rounds).sel(x=0, y=0),
)

day05.py

from re import compile as re_compile
from numpy import (
    array, zeros, put, linspace, dtype,
    ravel_multi_index, diff, hstack, repeat, arange,
)
from xarray import DataArray
from pandas import Categorical

LINE_RE = re_compile('\s*'.join([
    r'(\d+),(\d+)', r'->', r'(\d+),(\d+)',
]))
with open('data.txt') as f:
    entries = [
        [int(x) for x in mo.groups()]
        for mo in LINE_RE.finditer(f.read())
    ]

PointDtype = dtype([*{
    'x': int,
    'y': int,
}.items()])
points = array(entries).view(PointDtype)
runs = abs(hstack([
    diff(points['x']),
    diff(points['y']),
])) + 1
lengths = runs.max(axis=1)

lines = [
    repeat(arange(len(lengths)), lengths),
    hstack([
        linspace(*p['y'], l, dtype=int)
        for p, l in zip(points, lengths)
    ]),
    hstack([
        linspace(*p['x'], l, dtype=int)
        for p, l in zip(points, lengths)
    ]),
]

bounds = len(points), points['y'].max()+1, points['x'].max()+1
data = zeros(bounds, dtype=int)
indices = ravel_multi_index(lines, data.shape)
put(data, indices, 1)

floor = DataArray(
    data=data,
    dims=[*'zyx'],
    coords={
        'z': Categorical((runs > 1).all(axis=1))
             .rename_categories(['straight', 'diagonal']),
    },
)

print(
    (floor.sel(z='straight').sum(dim='z') > 1).sum(),
    (floor.sum(dim='z') > 1).sum(),
    sep='\n',
)

day06.py

from csv import reader
from pandas import Series, DataFrame
from numpy import arange, roll
from itertools import islice

def process(pop):
    while True:
        yield pop.copy()
        spawning, pop.loc[0] = pop.loc[0], 0
        pop.index = roll(pop.index, 1)
        pop.loc[6] += spawning
        pop.loc[8] += spawning

with open('data.txt') as f:
    init_pop = [int(x) for x in next(reader(f))]
    init_pop = Series(init_pop)
    init_pop = init_pop.value_counts().reindex(arange(9), fill_value=0)

populations = (
    DataFrame.from_records([
        pop for pop in islice(process(init_pop), 256+1)
    ])
    .rename_axis('day', axis='index')
    .rename_axis('lifetime', axis='columns')
)

print(
    populations.loc[80].sum(),
    populations.loc[256].sum(),
    sep='\n',
)

day07.py

from numpy import loadtxt, bincount, mgrid, diff, triu, tril, fliplr

positions = loadtxt('data.txt', delimiter=',', dtype=int)
positions = bincount(positions, minlength=positions.max())

costs = abs(diff(mgrid[:len(positions),:len(positions)], axis=0))[0]
new_costs = triu(costs).cumsum(axis=1) \
          + fliplr(fliplr(tril(costs)).cumsum(axis=1))

print(
    (positions @ costs).min(),
    (positions @ new_costs).min(),
    sep='\n',
)

day08.py

from numpy import frombuffer, unique, where, swapaxes, logspace, tile
from pandas import DataFrame, Series, MultiIndex, read_csv, merge
from scipy.ndimage import find_objects
from itertools import permutations, product

nums = b'''\
 aaaa        aaaa  aaaa       \
b    c     c     c     cb    c\
b    c     c     c     cb    c\
             dddd  dddd  dddd \
e    f     fe          f     f\
e    f     fe          f     f\
 gggg        gggg  gggg       \
 aaaa  aaaa  aaaa  aaaa  aaaa \
b     b          cb    cb    c\
b     b          cb    cb    c\
 dddd  dddd        dddd  dddd \
     fe    f     fe    f     f\
     fe    f     fe    f     f\
 gggg  gggg        gggg  gggg \
'''.replace(b'\\\n', b'')

raw_nums = frombuffer(nums, dtype='S1')
letters = unique(raw_nums[raw_nums != b' ']).astype('<U1')

nums = where(raw_nums == b' ', 0, raw_nums.view('int8') - ord(b'a') + 1)
nums = swapaxes(nums.reshape(2, 7, 5, 6), 1, 2).reshape(10, 7, 6)

digits = (
    DataFrame([
        Series(dict(zip(letters, find_objects(num, max_label=len(letters)))), name=val)
            .pipe(lambda s: ~s.isna())
        for val, num in enumerate(nums)
    ])
    .rename_axis('number').rename_axis('segment', axis='columns')
    .stack()
    .rename('active')
)

hashes = Series(
    logspace(len(letters) - 1, 0, len(letters), base=len(letters), dtype=int),
    index=letters,
)

all_orderings = Series(
    ([*order] for order in permutations(letters, r=len(letters)))
)

lookup = (
    all_orderings
    .explode()
    .pipe(lambda s:
        hashes.loc[s].set_axis(
            s.rename_axis('mapping').index.to_frame().assign(
                segment=tile(hashes.index, len(s) // len(hashes)),
            ).pipe(MultiIndex.from_frame)
        )
    )
    .pipe(lambda s: s * digits)
    .groupby(['mapping', 'number']).sum()
    .groupby(['mapping']).agg(frozenset)
    .rename('hash')
)

mappings = (
    all_orderings
    .apply(Series)
    .set_axis(letters, axis='columns')
    .stack()
    .rename_axis(['mapping', 'from_segment'])
    .rename('to_segment')
    .reset_index('from_segment')
    .set_index('to_segment', append=True)
    .squeeze()
    .rename_axis(['mapping', 'segment'])
    .rename('segment')
)

entries = (
    read_csv('data.txt',
        delimiter=' (?:\| )?',
        names=[
            *product(['control'], range(10)),
            *product(['output'], range(4))
        ],
        engine='python',
        header=None,
    ).rename_axis('entry')
)

controls = (
    entries['control']
    .rename_axis('code', axis='columns')
    .stack()
    .str.extractall(r'(?P<segment>\w)')
    .droplevel('match')
    .pipe(lambda df: df
        .assign(count=1)
        .pivot_table(
            index=['entry', 'code'],
            columns='segment',
            values='count',
            aggfunc='count',
            fill_value=0
        )
        .droplevel('code')
    )
    .pipe(lambda df: df @ hashes)
    .groupby('entry').agg(frozenset)
    .rename('hash')
)
outputs = (
    entries['output']
    .rename_axis('code', axis='columns')
    .stack()
    .str.extractall(r'(?P<segment>\w)')
    .droplevel('match')
)

decoded = (
    merge(
        lookup.pipe(lambda s:
            s.set_axis(
                MultiIndex.from_product([
                    [0], s.index
                ], names=['_', 'mapping'])
            ),
        ),
        controls.pipe(lambda s:
            s.set_axis(
                MultiIndex.from_product([
                    [0], s.index
                ], names=['_', 'entry'])
            ),
        ),
        left_index=True,
        right_index=True,
        suffixes=['_lookup', '_control'],
        how='outer'
    )
    .droplevel('_')
    .assign(match=lambda df: df['hash_lookup'] == df['hash_control'])
    .pipe(lambda df: df[df['match']])
    .pipe(lambda df: df.reset_index('mapping')['mapping'])
    .sort_index()
    .pipe(lambda s:
        outputs.join(s)
        .assign(active=True)
        .set_index(['mapping', 'segment'], append=True)
        .join(mappings)
        .groupby(['entry', 'code'])['segment'].agg(frozenset)
        .pipe(lambda s:
            digits[digits].reset_index('segment')['segment'].groupby('number').agg(frozenset)
            .reset_index('number')
            .set_index('segment')
            .loc[s]
            .set_axis(s.index)
        )
    )
    .squeeze()
)

print(
    decoded.groupby('entry').agg(
        lambda s: s.isin({1, 4, 7, 8}).sum()
    ).sum(),
    decoded.groupby('entry').agg(
        lambda s: s @ logspace(len(s) - 1, 0, len(s), base=10, dtype=int)
    ).sum(),
    sep='\n',
)

day09.py

from numpy import memmap, nonzero, arange, int8, pad, iinfo, newaxis, full, unique
from skimage.util import view_as_windows
from scipy.ndimage import generate_binary_structure, label
from pandas import Series

data = memmap('data.txt', mode='r', dtype='S1')
nls = nonzero(data == b'\n')[0]
width, height = nls[0], len(nls)
assert (nls == arange(nls[0], nls[-1]+1, width+1)).all()

ceiling = data.view(int8).reshape(height, width+1)[:, :width] - ord(b'0')
padded = pad(ceiling, ([1, 1], [1, 1]), constant_values=iinfo(ceiling.dtype).max)
neighbours = ~generate_binary_structure(2, 1)
neighbours[1, 1] = True
sections = view_as_windows(padded, neighbours.shape).reshape(width*height, *neighbours.shape)

sinks = (sections > (ceiling.ravel()[..., newaxis, newaxis])) | neighbours
sinks = sinks.all(axis=(1, 2)).reshape(ceiling.shape)
sink_coords = nonzero(sinks)

peak_coords = nonzero(ceiling == ceiling.max())
basins = full(fill_value=iinfo(ceiling.dtype).max, shape=ceiling.shape, dtype=int8)
basins[peak_coords] = 0
labels, sizes = unique(label(basins)[0], return_counts=True)
basins = Series(sizes, labels)

print(
    (ceiling[sink_coords] + 1).sum(),
    basins.loc[lambda s: s.index.difference({0})].nlargest(3).prod(),
    sep='\n',
)

day10.py

from pandas import Series, read_csv, DataFrame
from numpy import logspace

pairs = DataFrame([
    ('[', ']',),
    ('(', ')',),
    ('{', '}',),
    ('<', '>',),
]).set_axis(['opening', 'closing'], axis='columns')
opening = pairs.set_index('closing').squeeze()
closing = pairs.set_index('opening').squeeze()

error_scores = Series({')': 3, ']': 57, '}': 1197, '>': 25137,})
completion_scores = Series({')': 1, ']': 2, '}': 3, '>': 4,})

@lambda c: lambda *a, **kw: [ci := c(*a, **kw), next(ci), lambda val=None: ci.send(val)][-1]
def check():
    stack, value = [], None
    while True:
        valid = value not in opening or opening.loc[stack.pop()] == stack.pop()
        state = DataFrame({
            'valid':     valid,
            'size':      len(stack),
            'remaining': ''.join(closing.loc[stack].iloc[::-1]),
        }, index=[None])
        value = (yield state)['symbol'].item()
        stack.append(value)

entries = (
    read_csv(
        'data.txt',
        header=None,
        delimiter=None,
        squeeze=True,
    ).str.extractall(r'(?P<symbol>.)')
    .rename_axis(['entry', 'character'])
)

checked = (
    entries.pipe(lambda df:
        df.join(df
            .groupby('entry').apply(
                lambda g: g.groupby('character').apply(
                    check()
                ).droplevel(-1)
            )
        )
    )
)
erroneous = checked.groupby('entry')['valid'].filter(lambda g: not g.all()).index
incomplete = checked.groupby('entry')['size'].filter(lambda g: g.iloc[-1] > 0).index

print(
    checked
        .loc[erroneous]
        .groupby(['entry', 'character'])
        .filter(lambda g: ~g['valid'])
        .groupby(['entry']).first()
        ['symbol']
        .pipe(lambda s: error_scores.loc[s].set_axis(s.index))
        .sum()
    ,
    checked
        .loc[incomplete.difference(erroneous)]
        .groupby(['entry'])['remaining'].last()
        .str.extractall(r'(?P<symbol>.)')
        ['symbol']
        .groupby(['entry']).agg(
            lambda g: g
                .pipe(lambda s: completion_scores.loc[s].set_axis(s.index))
                .pipe(lambda s: s @ logspace(len(s) - 1, 0, len(s), dtype=int, base=5))
        )
        .median()
    ,
    sep='\n',
)

day11.py

from numpy import memmap, nonzero, arange, int8, array, meshgrid, zeros_like, unique, newaxis
from itertools import takewhile
from pandas import Series

def step(ὀκτώποδες):
    neighbours = array(meshgrid(offset := [-1, 0, +1], offset)).reshape(2, 9)
    all_flashed = zeros_like(ὀκτώποδες, dtype=bool)
    while True:
        ὀκτώποδες += 1
        all_flashed[:] = False
        while True:
            flashing = (ὀκτώποδες > 9) & ~all_flashed
            if not flashing.any():
                break
            affected = (xs := (
                array(nonzero(flashing))[:, newaxis, :]
              + neighbours[..., newaxis]
            ).reshape(2, -1))[:,
                (0 <= xs[0]) & (xs[0] < ὀκτώποδες.shape[0])
              & (0 <= xs[1]) & (xs[1] < ὀκτώποδες.shape[1])
            ]
            locs, vals = unique(affected, axis=1, return_counts=True)
            ὀκτώποδες[(*locs,)] += vals
            all_flashed |= flashing
        yield all_flashed
        ὀκτώποδες[ὀκτώποδες > 9] = 0

raw = memmap('data.txt', mode='r', dtype='S1')
nls, = nonzero(raw == b'\n')
width, height = nls[0], len(nls)
assert (nls == arange(nls[0], nls[-1] + 1, width + 1)).all()
ὀκτώποδες = raw.reshape(height, width + 1)[:, :-1].view(int8) - ord(b'0')

state = Series([
    *takewhile(
        lambda x: x != ὀκτώποδες.size,
        (x.sum() for x in step(ὀκτώποδες)),
    )
])

print(
    state.iloc[:100].sum(),
    len(state) + 1,
    sep='\n',
)

day12.py

from csv import reader
from networkx import DiGraph
from collections import namedtuple, Counter
from pandas import Series

class CaveSystem(DiGraph):
    def __init__(self, *args, **kwargs):
        self.nodes_by_name = {}
        super().__init__(*args, **kwargs)

    start = property(lambda s: s.nodes_by_name.get('start'))
    end   = property(lambda s: s.nodes_by_name.get('end'))
    def add_node_from_name(self, name):
        if name not in self.nodes_by_name:
            self.nodes_by_name[name] = Cave.from_name(name)
        return self.nodes_by_name[name]

    @classmethod
    def from_file(cls, filename):
        g = cls()
        with open(filename) as f:
            edges = {(u, v) for u, v in reader(f, delimiter='-')}
        edges = {
            (g.add_node_from_name(u), g.add_node_from_name(v))
            for u, v in edges
        }
        for u, v in edges:
            if u is not g.end and v is not g.start:
                g.add_edge(u, v)
            if v is not g.end and u is not g.start:
                g.add_edge(v, u)
        return g

    def traverse(self, *, predicate):
        start, end = self.start, self.end
        def traverse(node, visits):
            if not predicate(node, visits):
                return
            if node is end:
                yield visits
            for n in g.neighbors(node):
                yield from traverse(n, visits + Counter({node: 1}))
        yield from traverse(start, Counter())

class Cave(namedtuple('CaveBase', 'name')):
    SUBTYPES = {}
    @classmethod
    def from_name(cls, name):
        if cls not in cls.SUBTYPES:
            for subcls, (_, pred) in sorted(cls.SUBTYPES.items(), key=lambda kv: kv[-1][0]):
                if pred(name):
                    return subcls.from_name(name)
        return cls(name)
    def __init_subclass__(cls, predicate, *, priority=0):
        cls.SUBTYPES[cls] = priority, predicate
class StartCave(Cave, predicate=lambda name: name == 'start', priority=1): pass
class EndCave  (Cave, predicate=lambda name: name == 'end',   priority=1): pass
class SmallCave(Cave, predicate=lambda name: name.islower(),  priority=2): pass

g = CaveSystem.from_file('data.txt')

visit_once = lambda n, v: not isinstance(n, SmallCave) or v[n] < 1
visit_one_twice = lambda n, v: \
    not isinstance(n, SmallCave) \
    or v[n] < 1 \
    or sum(1 for x, c in v.items() if isinstance(x, SmallCave) and c > 1) < 1

paths1 = Series(g.traverse(predicate=visit_once), name='path')
paths2 = Series(g.traverse(predicate=visit_one_twice), name='path')

print(
    len(paths1),
    len(paths2),
    sep='\n',
)

day13.py

from itertools import takewhile
from numpy import (
    array, zeros, put, ravel_multi_index, pad, full, fliplr, flipud, where,
    array2string, swapaxes
)
from re import compile as re_compile
from dataclasses import dataclass
from collections import deque

class Fold:
    LINE_RE = re_compile(r'fold along (?P<axis>y|x)=(?P<index>\d+)')
    SUBTYPES = {}
    def __init_subclass__(cls, axis):
        cls.SUBTYPES[axis] = cls
    @classmethod
    def from_line(cls, line):
        mo = cls.LINE_RE.match(line)
        return cls.SUBTYPES[mo.group('axis')](int(mo.group('index')))
@dataclass
class Yfold(Fold, axis='y'):
    idx : int
    def __call__(self, data, *, fill_value=False):
        top, bot = data[:self.idx, :], data[self.idx+1:, :]
        sz = max(top.shape[0], bot.shape[0])
        top = pad(top, [(sz - top.shape[0], 0), (0, 0)], constant_values=fill_value)
        bot = pad(bot, [(0, sz - bot.shape[0]), (0, 0)], constant_values=fill_value)
        return (top | flipud(bot))[:sz, :]
@dataclass
class Xfold(Fold, axis='x'):
    idx : int
    def __call__(self, data, *, fill_value=False):
        lft, rgt = data[:, :self.idx], data[:, self.idx+1:]
        sz = max(lft.shape[1], rgt.shape[1])
        lft = pad(lft, [(0, 0), (sz - lft.shape[1], 0)], constant_values=fill_value)
        rgt = pad(rgt, [(0, 0), (0, sz - rgt.shape[1])], constant_values=fill_value)
        return (lft | fliplr(rgt))[:, :sz]

def fold(data, folds):
    for f in folds:
        yield (data := f(data))

with open('data.txt') as f:
    points = (ln.split(',') for ln in takewhile(lambda ln: ln.strip(), f))
    points = array([(int(y), int(x)) for x, y in points]).T
    folds = [Fold.from_line(ln) for ln in f]

height, width = points.max(axis=1) + 1
cave = zeros((height, width), dtype=bool)
put(cave, ravel_multi_index((*points,), (height, width)), True)

first_fold = next(fold(cave, folds))
last_fold = deque(fold(cave, folds), maxlen=1)[0]

print(
    first_fold.sum(),
   *(lambda letters: (
        array2string(lets, formatter={'bool': {True: '#', False: ' '}.get})
        for lets in letters
    ))(swapaxes(
        pad(xs := last_fold[:6, :40],
            [(0, 6 - xs.shape[0]), (0, 40 - xs.shape[1])],
            constant_values=False,
        ).reshape(6, -1, 4 * 5),
        1, 0
    )),
    sep='\n',
)

day14.py

from pandas import Series, MultiIndex, merge, DataFrame
from itertools import takewhile, islice, product
from re import compile as re_compile
from string import ascii_uppercase

def step(polymer, rules):
    while True:
        matches = (
            merge(rules, polymer, left_index=True, right_index=True)
            .pipe(lambda df: df.groupby(df['pairs'])['count'].sum())
            .pipe(lambda s: s.set_axis(MultiIndex.from_tuples(s.index, names=polymer.index.names)))
        )
        polymer = (
            polymer
          + matches.reindex(polymer.index, fill_value=0)
          - polymer.loc[rules.index.unique()].reindex(polymer.index, fill_value=0)
        )
        yield polymer

RULE_RE = re_compile(r'(?P<left>[A-Z])(?P<right>[A-Z]) -> (?P<middle>[A-Z])')
with open('data.txt') as f:
    polymer = ''.join(takewhile(bool, (ln.strip() for ln in f)))
    rules = (RULE_RE.match(ln) for ln in f)
    rules = Series({
        (mo.group('left'), mo.group('right')):
            [(mo.group('left'), mo.group('middle')),
             (mo.group('middle'), mo.group('right')),]
        for mo in rules
    }).rename_axis(['left', 'right']).rename('pairs').sort_index().explode()
    last = Series({polymer[-1]: 1}).reindex([*ascii_uppercase], fill_value=0)
    polymer = (
        Series(
            data=1,
            index=MultiIndex.from_tuples(zip(polymer, polymer[1:]), name=['left', 'right']),
            name='count',
        ).pipe(lambda s: s.groupby(['left', 'right']).sum())
        .reindex([*product(ascii_uppercase, repeat=2)], fill_value=0)
    )

state = DataFrame({
    st: pm
    for st, pm in
    enumerate(islice(step(polymer, rules), 40), start=1)
}).rename_axis('step', axis='columns')

print(
    state
    [[10, 40]]
    .pipe(lambda df: df.groupby('left').sum())
    .pipe(lambda df:
        df + DataFrame({col: last.loc[df.index] for col in df.columns})
    )
    .pipe(lambda df:
        Series({
            name: col[col>0]
               .sort_values()
               .pipe(lambda s: s.iloc[-1] - s.iloc[0])
            for name, col in df.items()
        }).rename_axis(df.columns.name)
    )
    ,
    sep='\n',
)

day15.py

from numpy import memmap, nonzero, arange, zeros, int8, kron, ones, mgrid
from scipy.sparse import csr_matrix, eye, kron as sparse_kron
from scipy.linalg import toeplitz
from scipy.sparse.csgraph import dijkstra

def cave_from_data(data):
    column = zeros(data.shape[0], dtype=int8)
    column[1] = 1
    offsets = csr_matrix(toeplitz(column), dtype=int8)
    ident = eye(offsets.shape[0], dtype=int8)
    return (
        (sparse_kron(offsets, ident) + sparse_kron(ident, offsets))
        .multiply(data.ravel())
    )

raw = memmap('data.txt', mode='r', dtype='S1')
nls, = nonzero(raw == b'\n')
width, height = nls[0], len(nls)
assert (nls == arange(nls[0], nls[-1] + 1, width + 1)).all()
data = raw.reshape(height, width + 1)[:, :width].astype(int8)

small_data = data
large_data = (
    kron(ones((sz := 5, sz), dtype=int8), data)
  + kron(mgrid[:sz, :sz].sum(axis=0), ones(data.shape, dtype=int8))
  - 1
) % 9 + 1

small_cave = cave_from_data(small_data)
large_cave = cave_from_data(large_data)

small_cave_weights = dijkstra(small_cave, directed=True, unweighted=False, indices=0)
large_cave_weights = dijkstra(large_cave, directed=True, unweighted=False, indices=0)

print(
    small_cave_weights[-1],
    large_cave_weights[-1],
    sep='\n',
)