🌹

toolkit.py

 def shortest_path(start, end, move):
     seen = {}
     edge = {start: None}
-    while end not in seen:
+    while edge:
         seen.update(edge)
         edge = {
             adj: pos
             for adj in move(pos)
             if adj not in seen
         }
-        if not edge:
-            raise RuntimeError('Path not found', seen)
+        if end in seen:
+            break
+    else:
+        raise RuntimeError('Path not found', seen)
     path = []
     while end:
         path.append(end)

y2019/p18.py

 import collections
 import itertools
 import sys
+from functools import partial
 
-from toolkit import read_image, render
+from toolkit import read_image, render, shortest_path
 
 
 def read_in(text):
     grid = read_image(text)
     elements = {v: k for k, v in grid.items()}
-    keys = frozenset(k for k in elements if 'a' <= k <= 'z')
+    keys = {c for c in elements if c.isalpha() and c.islower()}
     return dict(grid), elements, keys
 
 
-def move(pos):
+def move(grid, pos):
     for ori in [1, -1, 1j, -1j]:
         if grid[pos + ori] != '#':
             yield pos + ori
 
 
-def shortest_path(start, end, move):
-    seen = {}
-    edge = {start: None}
-    while end not in seen:
-        edge = {
-            adj: pos
-            for pos in edge
-            for adj in move(pos)
-            if adj not in seen
-        }
-        seen.update(edge)
-        if not edge:
-            raise RuntimeError('No path found')
-    path = []
-    while end != start:
-        path.append(end)
-        end = seen[end]
-    return path
-
-
-def trim(states):
+def trim(states, metric):
     groups = collections.defaultdict(list)
     for state in states:
         tips = frozenset(seq[-1] for seq in state)
         bulk = frozenset(' '.join(state))
         groups[bulk, tips].append(state)
-    return [min(groups[k], key=calc_total) for k in groups]
+    return [min(groups[k], key=metric) for k in groups]
 
 
-def precalc_moves(entrances):
+def precalc_moves(entrances, keys, elements, grid):
     requirements = collections.defaultdict(dict)
-    lengths = collections.defaultdict(dict)
-    for a, b in itertools.permutations(keys.union(entrances), 2):
+    lengths = {}
+    for a, b in itertools.combinations(keys.union(entrances), 2):
         try:
-            path = shortest_path(elements[a], elements[b], move)
+            path = shortest_path(elements[a], elements[b], partial(move, grid))
+            *glyphs, = map(grid.get, path)
         except RuntimeError:
             continue
-        lengths[a][b] = len(path)
-        requirements[a][b] = {
-            e.lower()
-            for e, v in elements.items()
-            if e.isupper()
-            and v in path
-        }
-    return requirements, lengths
+        lengths[a, b] = len(path) - 1
+        lengths[b, a] = len(path) - 1
+        requirements[a][b] = {c.lower() for c in glyphs if c.isupper()}
+        requirements[b][a] = {c.lower() for c in glyphs if c.isupper()}
+    return requirements, partial(calc_total, lengths)
 
 
-def calc_total(state):
-    return sum(sum(
-        lengths[a][b]
-        for a, b in zip(s, s[1:])
-        if {a, b} < set(lengths)
-    ) for s in state)
+def calc_total(lengths, state):
+    return sum(lengths[pair] for seq in state for pair in zip(seq, seq[1:]))
 
 
 def mod_grid(text):
     for char, ori in zip('*&^@', [1, -1, 1j, -1j]):
         grid[pos + ori] = '#'
         grid[pos + ori + ori * 1j] = char
-    return render(grid, {k: k for k in grid.values()})
+    return render(grid)
 
 
 def solve(text):
-    global grid, elements, keys, lengths
     grid, elements, keys = read_in(text)
     stacks = {k for k in elements if not k.isalnum()} - {'#', '.'}
-    requirements, lengths = precalc_moves(stacks)
+    requirements, metric = precalc_moves(stacks, keys, elements, grid)
     states = [stacks]
     for _ in range(len(keys)):
-        states = trim(
+        states = trim([
             state - {seq} | {seq + b}
             for state in states
             for seq in state
             for b, reqs in requirements[seq[-1]].items()
             if b not in seq and reqs.issubset(''.join(state))
-        )
-    final = min(states, key=calc_total)
-    print(' '.join(sorted(final)))
-    print(calc_total(final))
+        ], metric)
+    best = min(states, key=metric)
+    print(best)
+    print(metric(best))
 
 
 text = sys.stdin.read()