import collections import itertools import re MODES = { '0': lambda ns, x, relbase: ns[x], '1': lambda ns, x, relbase: x, '2': lambda ns, x, relbase: ns[x + relbase], '0-special': lambda ns, x, relbase: x, '2-special': lambda ns, x, relbase: x + relbase, } 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' yield pos + N for i in range(N): mode, x = modes[i], ns[pos + i] yield MODES[mode](ns, x, relbase) 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)[::-1] + '000') pos += 1 if op == 1: pos, a, b, c = consume(3, 'c') ns[c] = a + b elif op == 2: pos, a, b, c = consume(3, 'c') ns[c] = a * b elif op == 3: pos, a = consume(1, 'a') ns[a] = next(in_iter) elif op == 4: pos, a = consume(1) yield a elif op == 5: pos, a, b = consume(2) if a != 0: pos = b elif op == 6: pos, a, b = consume(2) if a == 0: pos = b elif op == 7: pos, a, b, c = consume(3, 'c') ns[c] = int(a < b) elif op == 8: pos, a, b, c = consume(3, 'c') ns[c] = int(a == b) elif op == 9: pos, a = consume(1) relbase += a elif op == 99: return else: raise RuntimeError(op) def parse(string): ns = [int(n) for n in re.findall(r'-?\d+', string)] memory = collections.defaultdict(int) memory.update(enumerate(ns)) return memory