day08.py

@@ -1,23 +1,9 @@
 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 = [
         [
@@ -35,21 +21,45 @@ class Day(AOCDay):
         inst = self.getInput()[0]
         nodes = {}
         for line in self.getInput()[2:]:
-            node, paths = line.split(" = ")
-            pl, pr = paths.split(", ")
+            node, pathhs = line.split(" = ")
+            pl, pr = pathhs.split(", ")
             nodes[node] = (pl[1:], pr[:-1])
 
         return inst, nodes
 
     def part1(self) -> Any:
         inst, nodes = self.parse_input()
-        return get_steps_to_z(inst, nodes, "AAA")
+        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
 
     def part2(self) -> Any:
         inst, nodes = self.parse_input()
+        i_len = len(inst)
         check_nodes = [x for x in nodes if x.endswith("A")]
-        z_steps = [get_steps_to_z(inst, nodes, x) for x in check_nodes]
-        return math.lcm(*z_steps)
+        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")])
 
 
 if __name__ == "__main__":