import collections import itertools import re def parse(string): ns = [int(n) for n in re.findall(r'-?\d+', string)] memory = collections.defaultdict(int) memory.update(enumerate(ns)) return memory def get_parameters(ns, pos, modes, N, writes, relbase): for c in writes: # paradox: return immediate mode to use positionally outside modes[ord(c) - ord('a')] += '-special' for mode, x in zip(modes, [ns[y] for y in range(pos, pos + N)]): yield { '0': lambda: ns[x], '1': lambda: x, '2': lambda: ns[x + relbase], '0-special': lambda: x, '2-special': lambda: x + relbase, }[mode]() yield pos + N def compute(ns, in_iter): def consume(N, writes=''): return get_parameters(ns, pos, modes, N, writes, relbase) if isinstance(ns, str): ns = parse(ns) if isinstance(in_iter, int): in_iter = itertools.cycle([in_iter]) pos = 0 relbase = 0 while True: op = ns[pos] % 100 # instructions stupidly say ABC referring to parameters 3, 2, 1 # we do a, b, c modes = list(str(ns[pos] // 100).zfill(3)[::-1]) pos += 1 if op == 1: a, b, c, pos = consume(3, 'c') ns[c] = a + b elif op == 2: a, b, c, pos = consume(3, 'c') ns[c] = a * b elif op == 3: a, pos = consume(1, 'a') ns[a] = next(in_iter) elif op == 4: a, pos = consume(1) yield a elif op == 5: a, b, pos = consume(2) if a != 0: pos = b elif op == 6: a, b, pos = consume(2) if a == 0: pos = b elif op == 7: a, b, c, pos = consume(3, 'c') ns[c] = int(a < b) elif op == 8: a, b, c, pos = consume(3, 'c') ns[c] = int(a == b) elif op == 9: a, pos = consume(1) relbase += a elif op == 99: return else: raise RuntimeError(op)