2024-22 python p2

1 files changed, 37 insertions, 3 deletions

diff --git a/2024/22/python/main.py b/2024/22/python/main.py
index 2354726..c290858 100644
--- a/2024/22/python/main.py
+++ b/2024/22/python/main.py
@@ -1,10 +1,11 @@
+from collections import Counter, deque
 from fileinput import input
 from itertools import islice
 
 lines = [int(line.strip()) for line in input()]
 
 
-def prices(secret):
+def secrets(secret):
     x = secret
     while True:
         yield x
@@ -16,8 +17,41 @@ def prices(secret):
         x &= 0xFFFFFF
 
 
-silver = sum(next(islice(prices(x), 2000, None)) for x in lines)
-gold = 0
+def prices(it):
+    for s in it:
+        yield s % 10
+
+
+def diff(it):
+    prev = next(it)
+    for curr in it:
+        yield curr, curr - prev
+        prev = curr
+
+
+def to_seqs(it, n=4):
+    head = map(lambda x: x[1], islice(it, n - 1))
+    q = deque(head, maxlen=n)
+    for price, diff in it:
+        q.append(diff)
+        yield price, tuple(q)
+
+
+def seq_sells(it):
+    sells = {}
+    for price, diff in it:
+        if diff in sells:
+            continue
+        sells[diff] = price
+    return sells
+
+
+best_seqs = Counter()
+for x in lines:
+    best_seqs += seq_sells(to_seqs(islice(diff(prices(secrets(x))), 2000)))
+
+silver = sum(next(islice(secrets(x), 2000, None)) for x in lines)
+gold = best_seqs.most_common(1)[0][1]
 
 print("silver:", silver)
 print("gold:", gold)