Ver código fonte

Simplify day 21

Taddeus Kroes 1 ano atrás
pai
commit
2f57606af5
1 arquivos alterados com 10 adições e 11 exclusões
  1. 10 11
      2024/21_keypad.py

+ 10 - 11
2024/21_keypad.py

@@ -3,24 +3,23 @@ import sys
 from functools import cache
 from functools import cache
 from itertools import pairwise
 from itertools import pairwise
 
 
-NUMERIC = '789', '456', '123', ' 0A'
-DIRECTIONAL = ' ^A', '<v>'
+NUMERIC = '789' '456' '123' ' 0A'
+DIRECTIONAL = ' ^A' '<v>'
 
 
 def walk(keypad, x, y, path):
 def walk(keypad, x, y, path):
+    i = y * 3 + x
     for direction in path:
     for direction in path:
-        neighbors = (x - 1, y), (x + 1, y), (x, y - 1), (x, y + 1)
-        x, y = neighbors['<>^v'.index(direction)]
-        yield keypad[y][x]
+        i += (-1, 1, -3, 3)['<>^v'.index(direction)]
+        yield keypad[i]
 
 
 def paths_between(keypad, start, end):
 def paths_between(keypad, start, end):
-    x1, y1 = next((x, y) for y, row in enumerate(keypad)
-                         for x, key in enumerate(row) if key == start)
-    x2, y2 = next((x, y) for y, row in enumerate(keypad)
-                         for x, key in enumerate(row) if key == end)
+    y1, x1 = divmod(keypad.index(start), 3)
+    y2, x2 = divmod(keypad.index(end), 3)
     hor = '<>'[x2 > x1] * abs(x2 - x1)
     hor = '<>'[x2 > x1] * abs(x2 - x1)
     ver = '^v'[y2 > y1] * abs(y2 - y1)
     ver = '^v'[y2 > y1] * abs(y2 - y1)
-    return tuple(path + 'A' for path in {hor + ver, ver + hor}
-                 if ' ' not in walk(keypad, x1, y1, path))
+    for path in {hor + ver, ver + hor}:
+        if ' ' not in walk(keypad, x1, y1, path):
+            yield path + 'A'
 
 
 @cache
 @cache
 def cost_between(keypad, start, end, links):
 def cost_between(keypad, start, end, links):