import sys import re from math import gcd from collections import defaultdict, Counter D = open(sys.argv[1]).read().strip() L = D.split("\n") def lcm(xs): ans = 1 for x in xs: ans = (x * ans) // gcd(x, ans) return ans GO = {"L": {}, "R": {}} steps, rule = D.split("\n\n") for line in rule.split("\n"): st, lr = line.split("=") st = st.strip() left, right = lr.split(",") left = left.strip()[1:].strip() right = right[:-1].strip() GO["L"][st] = left GO["R"][st] = right def solve(part2): POS = [] for s in GO["L"]: if s.endswith("A" if part2 else "AAA"): POS.append(s) T = {} t = 0 while True: NP = [] for i, p in enumerate(POS): p = GO[steps[t % len(steps)]][p] if p.endswith("Z"): T[i] = t + 1 # we need to solve t % A[i] == 0 for 6 different values of A[i] # in general, you could have to sole t % A[i] = B[i]. # you can do this with the chinese remainder theorem. if len(T) == len(POS): return lcm(T.values()) NP.append(p) POS = NP t += 1 assert False print(solve(False)) print(solve(True))