aoc/2019/22_spacecards.py

#!/usr/bin/env python3
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):
        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(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 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 frepeat(f, n, mod):
    if n == 0:
        return 1, 0
    if n == 1:
        return f
    half, odd = divmod(n, 2)
    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):
    f = inverse_function(instructions, ncards)
    fn = frepeat(f, nshuffles, ncards)
    return fapply(fn, index, ncards)

instructions = list(sys.stdin)
print(find_card(2019, 10007, instructions))
print(card_at(2020, 119315717514047, 101741582076661, instructions))