Cleanup

2020/08_loop.py

@@ -29,9 +29,9 @@ def patch(program):
         if opcode in change:
             program[i] = change[opcode], arg
             term, acc = run(program)
-            program[i] = opcode, arg
             if term:
                 return acc
+            program[i] = opcode, arg
 
 program = list(parse(sys.stdin))
 print(run(program)[1])

2020/14_docking.py

@@ -15,27 +15,32 @@ def parse(f):
         else:
             yield False, int(left[4:-1]), int(right)
 
-def run(program, v2):
+def run_v1(program):
     mem = {}
     and_mask = or_mask = 0
-    fluct = []
     for is_mask, a, b in program:
-        if v2 and is_mask:
-            and_mask = ~a & ~b
-            or_mask = b
-            fluct = [i for i in range(36) if (a >> i) & 1]
-        elif v2:
-            for bits in product((0, 1), repeat=len(fluct)):
-                fluct_bits = (bit << i for bit, i in zip(bits, fluct))
-                fluct_mask = reduce(or_, fluct_bits, 0)
-                mem[a & and_mask | or_mask | fluct_mask] = b
-        elif is_mask:
+        if is_mask:
             and_mask = a
             or_mask = b
         else:
             mem[a] = b & and_mask | or_mask
     return sum(mem.values())
 
+def run_v2(program):
+    mem = {}
+    fluct_masks = []
+    for is_mask, a, b in program:
+        if is_mask:
+            and_mask = ~(a | b)
+            or_mask = b
+            fl = [i for i in range(36) if a >> i & 1]
+            fluct_masks = [reduce(or_, (b << i for b, i in zip(bits, fl)), 0)
+                           for bits in product((0, 1), repeat=len(fl))]
+        else:
+            for fluct_mask in fluct_masks:
+                mem[a & and_mask | or_mask | fluct_mask] = b
+    return sum(mem.values())
+
 program = list(parse(sys.stdin))
-print(run(program, False))
-print(run(program, True))
+print(run_v1(program))
+print(run_v2(program))