Cleanup

2019/22_spacecards.py

@@ -2,42 +2,58 @@
 import sys
 from functools import partial, reduce
 
+# get modular multiplicative inverse for prime p
+def prime_modinv(p, mod):
+    return pow(p, mod - 2, mod)
+
 # get (a, b) coefficients for inverse linear index mapping functions ax+b
 def inverse_functions(instructions, ncards):
     for line in reversed(instructions):
-        words = line.split()
-        if words[0] == 'cut':
-            yield 1, int(words[1])
-        elif words[1] == 'into':
+        if line == 'deal into new stack\n':
             yield -1, ncards - 1
+        elif line.startswith('cut'):
+            amount = int(line.split()[1])
+            yield 1, amount
         else:
-            increment = int(words[-1])
-            index_increment = pow(increment, ncards - 2, ncards)
-            yield index_increment, 0
+            increment = int(line.split()[-1])
+            yield prime_modinv(increment, ncards), 0
+
+# compose all inverse mappings into a single linear function
+def inverse_function(instructions, ncards):
+    functions = inverse_functions(instructions, ncards)
+    return reduce(partial(fcompose, ncards), functions, (1, 0))
 
 # compose linear functions ax+b and cx+d
-def compose(mod, f, g):
+def fcompose(mod, f, g):
     a, b = f
     c, d = g
     return c * a % mod, (c * b + d) % mod
 
+# compute f(x) = ax+b
+def fapply(f, x, mod):
+    a, b = f
+    return (a * x + b) % mod
+
 # repeat ax+b n times
-def repeat_linear(f, n, mod):
+def frepeat(f, n, mod):
     if n == 0:
         return 1, 0
     if n == 1:
         return f
     half, odd = divmod(n, 2)
-    g = repeat_linear(f, half, mod)
-    g = compose(mod, g, g)
-    return compose(mod, f, g) if odd else g
+    g = frepeat(f, half, mod)
+    gg = fcompose(mod, g, g)
+    return fcompose(mod, f, gg) if odd else gg
+
+def find_card(value, ncards, instructions):
+    f = inverse_function(instructions, ncards)
+    return next(i for i in range(ncards) if fapply(f, i, ncards) == value)
 
 def card_at(index, ncards, nshuffles, instructions):
-    functions = inverse_functions(instructions, ncards)
-    invmap = reduce(partial(compose, ncards), functions, (1, 0))
-    a, b = repeat_linear(invmap, nshuffles, ncards)
-    return (index * a + b) % ncards
+    f = inverse_function(instructions, ncards)
+    fn = frepeat(f, nshuffles, ncards)
+    return fapply(fn, index, ncards)
 
 instructions = list(sys.stdin)
-print(next(i for i in range(10007) if card_at(i, 10007, 1, instructions) == 2019))
+print(find_card(2019, 10007, instructions))
 print(card_at(2020, 119315717514047, 101741582076661, instructions))