Day 8 - cleanup

day08.py

@@ -1,9 +1,23 @@
 import math
-from collections import defaultdict
 from tools.aoc import AOCDay
 from typing import Any
 
 
+def get_steps_to_z(inst: str, nodes: dict, from_node: str) -> int:
+    i_len = len(inst)
+    cur_step = 0
+    cur_node = from_node
+    while not cur_node.endswith("Z"):
+        if inst[cur_step % i_len] == "L":
+            cur_node = nodes[cur_node][0]
+        else:
+            cur_node = nodes[cur_node][1]
+
+        cur_step += 1
+
+    return cur_step
+
+
 class Day(AOCDay):
     inputs = [
         [
@@ -21,45 +35,21 @@ class Day(AOCDay):
         inst = self.getInput()[0]
         nodes = {}
         for line in self.getInput()[2:]:
-            node, pathhs = line.split(" = ")
-            pl, pr = pathhs.split(", ")
+            node, paths = line.split(" = ")
+            pl, pr = paths.split(", ")
             nodes[node] = (pl[1:], pr[:-1])
 
         return inst, nodes
 
     def part1(self) -> Any:
         inst, nodes = self.parse_input()
-        cur_node = "AAA"
-        cur_step = 0
-        while cur_node != "ZZZ":
-            dir = inst[cur_step % len(inst)]
-            if dir == "L":
-                cur_node = nodes[cur_node][0]
-            else:
-                cur_node = nodes[cur_node][1]
-
-            cur_step += 1
-
-        return cur_step
+        return get_steps_to_z(inst, nodes, "AAA")
 
     def part2(self) -> Any:
         inst, nodes = self.parse_input()
-        i_len = len(inst)
         check_nodes = [x for x in nodes if x.endswith("A")]
-        nodes_dp = defaultdict(dict)
-        for node in check_nodes:
-            cur_step = 0
-            cur_node = node
-            while (cur_node, cur_step % i_len) not in nodes_dp[node]:
-                nodes_dp[node][(cur_node, cur_step % i_len)] = cur_step
-                if inst[cur_step % i_len] == "L":
-                    cur_node = nodes[cur_node][0]
-                else:
-                    cur_node = nodes[cur_node][1]
-
-                cur_step += 1
-
-        return math.lcm(*[v for node in check_nodes for k, v in nodes_dp[node].items() if k[0].endswith("Z")])
+        z_steps = [get_steps_to_z(inst, nodes, x) for x in check_nodes]
+        return math.lcm(*z_steps)
 
 
 if __name__ == "__main__":