Advent of Code 2024, Day 21

This one was confusing.

My first idea was to transform each string of button presses into the next string, in its entirety, and do that for each level. It worked out for part 1. For part 2, it was too slow.

Then I found a tip: focus on the button level (instead of on the sequence level) and return numbers (instead of strings); have a function that takes in two buttons and the indirection number and use memoization there. I applied the idea to my code and it worked. I had already thought about memoization for the breadth-first search, but it was of no use there. My code now runs in less than 25 milliseconds on my machine, for both parts.

During the implementation, I got confused about when to prepend and when to append “A” button presses.

• In order to realize a sequence, you have to move the robot from “A” (the starting position) to the first button and push it, then move it from the first to the second button and push it, and so on. That is why you have to prepend an “A” there.
• In order to move the robot from one button to another and push it, you have to push the correct arrow buttons one level behind it (to move the robot) and then push the “A” button there (to actually push the button you have moved to). That is why you have to append an “A” there.

For the breadth-first search, I added a border of “#” signs to my maps so that I wouldn’t have to do any numerical boundary checks on the coordinates, just like in the puzzle inputs of previous days.

Both parts

For part 1, give “2” as the second argument. For part 2, give “25” as the second argument.

import sys

numpad = [
    "#####",
    "#789#",
    "#456#",
    "#123#",
    "##0A#",
    "#####"
]

dirpad = [
    "#####",
    "##^A#",
    "##",
    "#####"
]

def arrow_paths_to(i, j, pred):
    if (i, j) not in pred: return [""]
    paths = []
    arrows = {(-1, 0): "^", (1, 0): "v", (0, -1): "<", (0, 1): ">"}
    for oi, oj in pred[(i, j)]:
        last_arrow = arrows[(i-oi, j-oj)]
        for path in arrow_paths_to(oi, oj, pred):
            paths.append(path + last_arrow)
    return paths

def bfs_pred(pad, si, sj):
    q = [(si, sj)]
    d = {(si, sj): 0}
    pred = {}
    while len(q) > 0:
        i, j = q.pop(0)
        for ni, nj in [(i-1, j), (i+1, j), (i, j-1), (i, j+1)]:
            if pad[ni][nj] == "#": continue
            if (ni, nj) not in d:
                q.append((ni, nj))
                d[(ni, nj)] = d[(i, j)]+1
                pred[(ni, nj)] = {(i, j)}
            elif d[(ni, nj)] == d[(i, j)]+1:
                pred[(ni, nj)].add((i, j))
    return pred

def arrow_paths_between(a, b):
    if a == b: return [""]
    pad = numpad if a.isdigit() or b.isdigit() else dirpad
    for i, line in enumerate(pad):
        for j, c in enumerate(line):
            if c == a: si, sj = i, j
            if c == b: ei, ej = i, j
    return arrow_paths_to(ei, ej, bfs_pred(pad, si, sj))

mem = {}
def buttons_between(a, b, level):
    if (a, b, level) in mem: return mem[(a, b, level)]
    paths = arrow_paths_between(a, b)
    seqs = [path + "A" for path in paths]
    r = min([buttons_to_push(seq, level) for seq in seqs])
    mem[(a, b, level)] = r
    return r

def buttons_to_push(seq, level):
    if level == 0: return len(seq)
    seq = "A" + seq
    s = 0
    for i in range(1, len(seq)):
        a, b = seq[i-1], seq[i]
        s += buttons_between(a, b, level-1)
    return s

with open(sys.argv[1]) as f:
    codes = [line.strip() for line in f]

levels = int(sys.argv[2])

s = 0
for code in codes:
    n = buttons_to_push(code, levels+1)
    x = int("0" + code.replace("A", ""))
    s += n*x
print(s)

2024-12-22