2020 day 20-22

2020/d22.py

@@ -0,0 +1,86 @@
+from collections import deque, defaultdict
+from lib import get_data, ints
+
+data = get_data(__file__)
+
+a, b = data.strip().split("\n\n")
+a = deque(ints(a)[1:])
+b = deque(ints(b)[1:])
+
+while a and b:
+    x, y = a.popleft(), b.popleft()
+    if x > y:
+        a.append(x)
+        a.append(y)
+    elif y > x:
+        b.append(y)
+        b.append(x)
+    else:
+        assert False
+
+
+for p in [a, b]:
+    s = 0
+    for i, c in enumerate(reversed(p)):
+        s += (i + 1) * c
+    if s > 0:
+        print(s)
+
+
+CACHE = {}
+
+
+def play(a, b):
+    if len(a) == 0 or len(b) == 0:
+        return a, b
+
+    s = (tuple(a), tuple(b))
+    if s in CACHE:
+        return CACHE[s]
+
+    seen = set()
+    while a and b:
+        s = (tuple(a), tuple(b))
+        if s in seen:
+            return True, False
+        seen.add(s)
+
+        x, y = a.popleft(), b.popleft()
+        if x <= len(a) and y <= len(b):
+            ar, br = play(deque(list(a)[:x]), deque(list(b)[:y]))
+            if ar:
+                a.append(x)
+                a.append(y)
+            elif br:
+                b.append(y)
+                b.append(x)
+            else:
+                assert False
+        else:
+            if x > y:
+                a.append(x)
+                a.append(y)
+            elif y > x:
+                b.append(y)
+                b.append(x)
+            else:
+                assert False
+
+    s = (tuple(a), tuple(b))
+    if s not in CACHE:
+        CACHE[s] = (a, b)
+
+    return a, b
+
+
+a, b = data.strip().split("\n\n")
+a = deque(ints(a)[1:])
+b = deque(ints(b)[1:])
+a, b = play(a, b)
+
+for p in [a, b]:
+    s = 0
+    for i, c in enumerate(reversed(p)):
+        s += (i + 1) * c
+    if s > 0:
+        print(s)