from functools import cache DKP = { "^": (0, 1), "A": (0, 2), "<": (1, 0), "v": (1, 1), ">": (1, 2), } NKP = { "7": (0, 0), "8": (0, 1), "9": (0, 2), "4": (1, 0), "5": (1, 1), "6": (1, 2), "1": (2, 0), "2": (2, 1), "3": (2, 2), "0": (3, 1), "A": (3, 2), } def solve(data): codes = data.split() p1 = 0 p2 = 0 for code in codes: num = int(code[:-1]) p1 += num * unfold(code, 1 + 2, num=True) p2 += num * unfold(code, 1 + 25, num=True) return p1, p2 @cache def unfold(sequence, iterations, num=False): if iterations == 0: return len(sequence) s = "A" total = 0 pad, invalid = (NKP, (3, 0)) if num else (DKP, (0, 0)) for d in sequence: total += unfold(seq(pad[s], pad[d], invalid), iterations - 1) s = d return total @cache def seq(S, E, invalid): y1, x1 = S y2, x2 = E seq = "<" * (x1 - x2) + "v" * (y2 - y1) + "^" * (y1 - y2) + ">" * (x2 - x1) if (y2, x1) == invalid or (y1, x2) == invalid: seq = seq[::-1] return seq + "A" if __name__ == "__main__": with open("./input/21.txt", "r") as f: inp = f.read().strip() p1, p2 = solve(inp) print(p1) print(p2)