2024-22 python change to generator

1 files changed, 6 insertions, 4 deletions

diff --git a/2024/22/python/main.py b/2024/22/python/main.py
index 354ae78..2354726 100644
--- a/2024/22/python/main.py
+++ b/2024/22/python/main.py
@@ -1,20 +1,22 @@
 from fileinput import input
+from itertools import islice
 
 lines = [int(line.strip()) for line in input()]
 
 
-def evolve(x, n):
-    for _ in range(n):
+def prices(secret):
+    x = secret
+    while True:
+        yield x
         x ^= x << 6
         x &= 0xFFFFFF
         x ^= x >> 5
         x &= 0xFFFFFF
         x ^= x << 11
         x &= 0xFFFFFF
-    return x
 
 
-silver = sum(evolve(x, 2000) for x in lines)
+silver = sum(next(islice(prices(x), 2000, None)) for x in lines)
 gold = 0
 
 print("silver:", silver)