Part two

2024/day11/main.py

@@ -4,7 +4,9 @@ import os
 sys.path.insert(1, os.path.abspath('../../'))
 from python_tools.aoc_utils import *
 
-G = {}
+import collections as C
+
+G = C.defaultdict(lambda: 0)
 
 def blink_rocks(r):
     if r == 0:
@@ -15,47 +17,30 @@ def blink_rocks(r):
     else:
         return [r*2024]
 
-def blink_rocksp2(r):
-    if r not in G:
-        x = [r]
-        for i in range(3):
+
+def solve(stones, num):
+    G.clear()
+    for i in stones:
+        G[i] += 1
+
+    for i in range(num):
         nx = []
-            for j in x:
-                nx.extend(blink_rocks(j))
-            x = nx
-        G[r] = x
-    return G[r]
-
-def part_one(stones):
-    x = stones
-    for i in range(25):
-        nx = []
-        for j in x:
-            nx.extend(blink_rocks(j))
-        x = nx
-    print("Part One: ", len(x))
-
-STONES = set()
-
-def part_two(stones):
-    x = stones
-    for i in range(25):
-        print(i)
-        nx = []
-        for j in x:
-            nx.extend(blink_rocks(j))
-        for i in nx:
-            STONES.add(i)
-        x = nx
-        print("Len_stones: ", len(STONES))
-
-    print("Part Two: ", len(x))
+        g = G.copy()
+        G.clear()
+        for j in g:
+            result = blink_rocks(j)
+            for k in result:
+                G[k]+=g[j]
+    ans = 0
+    for i in G:
+        ans += G[i]
+    print(f"Stones after {num} blinks: {ans}")
 
 def main():
     f = open("input.txt", "r")
     contents = list(map(int, f.read().strip().split()))
-    part_one(contents)
-    part_two(contents)
+    solve(contents, 25)
+    solve(contents, 75)
 
 if __name__ == "__main__":
     main()