day 11 python cleanup

day11/sol.py

@@ -1,54 +1,71 @@
-from collections import Counter
+from collections import Counter, namedtuple
+import math
+
+Monkey = namedtuple('Monkey', 'index, op, divisor, target')
 def parse(input):
-    if hasattr(input, 'read'):
-        input = input.read()
     monkeys = []
-    for chunk in input.split("\n\n"):
+    items = {}
+    for i, chunk in enumerate(input.split("\n\n")):
         lines = chunk.strip().split("\n")
-        assert 'Starting' in lines[1]
-        items = [int(x.strip()) for x in lines[1].split(":")[1].split(",")]
-        assert 'Operation' in lines[2]
+        assert 'Starting items:' in lines[1]
+        assert 'Operation: new =' in lines[2]
+        assert 'Test: divisible by' in lines[3]
+        assert 'If true: throw to' in lines[4]
+        assert 'If false: throw to' in lines[5]
+
+        items[i] = [int(x.strip()) for x in lines[1].split(":")[1].split(",")]
         op = eval("lambda old: " + lines[2].split("=")[1])
-        assert 'divisible by' in lines[3]
         divisor = int(lines[3].split()[-1])
-        assert 'true: throw' in lines[4]
-        assert 'false: throw' in lines[5]
         target = [
             int(lines[5].split()[-1]),
             int(lines[4].split()[-1]),
         ]
+        monkeys.append(Monkey(items, op, divisor, target))
+    return monkeys, items
 
-        monkeys.append((items, op, divisor, target))
-    return monkeys
-
-def product(it):
-    t = 1
-    for x in it:
-        t  *= x
-    return t
-
-throws = Counter()
-def play(monkeys):
-    N = product(m[2] for m in monkeys)
-    for i, m in enumerate(monkeys):
-        items, op, divisor, target = m
-        for j, x in enumerate(items):
-            throws[i] += 1
-            old = x
-            x = op(x)
-            #x //= 3
-            x %= N
-            trg = target[x % divisor == 0]
-            #print(f'{old} -> {x}. throwing to {trg}')
-            monkeys[trg][0].append(x)
-        items[:] = []
+def play(monkeys, items, rounds=1, N=None):
+    throws = Counter()
+    for _ in range(rounds):
+        for i, m in enumerate(monkeys):
+            for j, x in enumerate(items[i]):
+                if N:
+                    nx = m.op(x) % N
+                else:
+                    nx = m.op(x) // 3
+                trg = m.target[nx % m.divisor == 0]
+                #print(f'{i}: {x} -> {nx}. throwing to {trg}')
+                items[trg].append(nx)
+            throws[i] += len(items[i])
+            items[i] = []
     return throws
 
-m = parse(open("input"))
-print(m)
-for _ in range(10000):
-    play(m)
-for i, m in enumerate(m):
-    print(i, m[0])
-a, b = sorted(throws.values())[-2:]
-print(a*b)
+def show(items):
+    for k in items:
+        print(k, items[k])
+
+def monkeybusiness(throws):
+    a, b = [x[1] for x in throws.most_common(2)]
+    return a*b
+
+def lcm(ints):
+    n = 1
+    for x in ints:
+        n *= x // math.gcd(n, x)
+    return n
+
+def main():
+    monkeys, items = parse(open("input").read())
+    #print(monkeys)
+    show(items)
+
+    tmp = {i: list(xs) for i, xs in items.items()}
+    throws = play(monkeys, tmp, rounds=20)
+    show(tmp)
+    print(monkeybusiness(throws))
+
+    N = lcm(m.divisor for m in monkeys)
+    throws = play(monkeys, items, rounds=10000, N=N)
+    show(items)
+    print(monkeybusiness(throws))
+
+main()