advent-of-code/2024-python/output/day_21.py

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)