aoc2024day21.py

# === Parsing ===

with open("input.txt", "r") as f:
    codes = [line.strip() for line in f]

print(codes)


# === Part 1 ===
numpad_pos = {
    "7": (0, 0),
    "8": (1, 0),
    "9": (2, 0),
    "4": (0, 1),
    "5": (1, 1),
    "6": (2, 1),
    "1": (0, 2),
    "2": (1, 2),
    "3": (2, 2),
    "0": (1, 3),
    "A": (2, 3),
}

arrowpad_pos = {
    "^": (1, 0),
    "A": (2, 0),
    "<": (0, 1),
    "v": (1, 1),
    ">": (2, 1),
}


def build_translation_mapping(pos):
    movements = {}
    for from_, (fx, fy) in pos.items():
        for to_, (tx, ty) in pos.items():
            if tx == 0 and fy == 3:
                movements[from_, to_] = "^" * (fy - ty) + "<" * (fx - tx)
            elif fx == 0 and ty == 3:
                movements[from_, to_] = ">" * (tx - fx) + "v" * (ty - fy)
            else:
                movements[from_, to_] = (
                    "<" * max(0, fx - tx)
                    + "^" * max(0, fy - ty)
                    + "v" * max(0, ty - fy)
                    + ">" * max(0, tx - fx)
                )

    return movements


numpad_movements = build_translation_mapping(numpad_pos)

arrowpad_movements = {
    ("^", "^"): "",
    ("^", "A"): ">",
    ("^", "<"): "v<",
    ("^", "v"): "v",
    ("^", ">"): "v>",
    ("A", "^"): "<",
    ("A", "A"): "",
    ("A", "<"): "v<<",
    ("A", "v"): "<v",
    ("A", ">"): "v",
    ("<", "^"): ">^",
    ("<", "A"): ">>^",
    ("<", "<"): "",
    ("<", "v"): ">",
    ("<", ">"): ">>",
    ("v", "^"): "^",
    ("v", "A"): "^>",
    ("v", "<"): "<",
    ("v", "v"): "",
    ("v", ">"): ">",
    (">", "^"): "<^",
    (">", "A"): "^",
    (">", "<"): "<<",
    (">", "v"): "<",
    (">", ">"): "",
}

from itertools import pairwise


def translate_sequence(sequence, movements):
    final = []
    for from_, to_ in pairwise("A" + sequence):
        final.append(movements[from_, to_] + "A")
    return "".join(final)


total_complexity = 0
for code in codes:
    lvl1 = translate_sequence(code, numpad_movements)
    lvl2 = translate_sequence(lvl1, arrowpad_movements)
    lvl3 = translate_sequence(lvl2, arrowpad_movements)
    total_complexity += len(lvl3) * int(code[:3])

print(total_complexity)


# === Part 2 ===
def to_segments(sequence):
    segments = []
    ptr = 0
    while ptr < len(sequence):
        seg_start = ptr
        seg_end = sequence.index("A", seg_start)
        while seg_end < len(sequence) and sequence[seg_end] == "A":
            seg_end += 1
        segments.append(sequence[seg_start:seg_end])
        ptr = seg_end
    return tuple(segments)


def translate_segment(segment, movements):
    segments = to_segments(translate_sequence(segment, movements))
    return segments


from collections import Counter


total_complexity = 0
for code in codes:
    as_arrows = translate_sequence(code, numpad_movements)
    segments = to_segments(as_arrows)
    ctr = Counter(segments)
    for _ in range(25):
        new_ctr = Counter()
        for segment, count in ctr.items():
            translation = translate_segment(segment, arrowpad_movements)
            for seg in translation:
                new_ctr[seg] += count
        ctr = new_ctr

    length = sum(len(segment) * count for segment, count in ctr.items())
    total_complexity += length * int(code[:3])

print(total_complexity)  # 167_538_833_832_712