Advent of Code day 23 part 2

AoC_2021_23_2.py

import networkx as nx


# indexing of positions is as follows:
#  0  1  2  3  4  5  6  7  8  9  10
#       11    12    13    14
#       15    16    17    18
#       19    20    21    22
#       23    24    25    26
ROOMS = {
    "A": {11, 15, 19, 23},
    "B": {12, 16, 20, 24},
    "C": {13, 17, 21, 25},
    "D": {14, 18, 22, 26},
}
# fmt: off
NEIGHBRS = [(0, 1), (1, 0), (1, 2), (2, 1), (2, 3), (2, 11), (3, 2), (3, 4), (4, 3), (4, 5), (4, 12), (5, 4), (5, 6), (6, 5), (6, 7), (6, 13), (7, 6), (7, 8), (8, 7), (8, 9), (8, 14), (9, 8), (9, 10), (10, 9), (11, 2), (11, 15), (12, 4), (12, 16), (13, 6), (13, 17), (14, 8), (14, 18), (15, 11), (15, 19), (16, 12), (16, 20), (17, 13), (17, 21), (18, 14), (18, 22), (19, 15), (19, 23), (20, 16), (20, 24), (21, 17), (21, 25), (22, 18), (22, 26), (23, 19), (24, 20), (25, 21), (26, 22)]  # noqa
# fmt: on
G = nx.Graph()
G.add_edges_from(NEIGHBRS)
DISTANCE = dict(nx.shortest_path_length(G))
SHORTEST_PATHS = [
    path
    for frm, paths in nx.shortest_path(G).items()
    for to_, path in paths.items()
    if len(path) > 1
    and to_ not in {2, 4, 6, 8}  # no door blocking
    and not {frm, to_}.issubset(set(range(11)))  # no hall-to-hall paths
]


def travel(fld, path):
    frm, to_ = path[0], path[-1]
    chars, guests = fld["chars"], fld["guests"]
    new_e = fld["spent"] + 10 ** (ord(chars[frm]) - ord("A")) * (len(path) - 1)
    a, b = min(frm, to_), max(frm, to_)
    new_c = chars[:a] + chars[b] + chars[a + 1 : b] + chars[a] + chars[b + 1 :]
    new_i = f"{new_e:05d}_{new_c}"
    new_g = {ch: guests[ch] - {frm} for ch in "ABCD"}
    return {"ident": new_i, "spent": new_e, "chars": new_c, "guests": new_g}


def min_future_spent(fld):

    ch3, ch5, ch7 = fld["chars"][3], fld["chars"][5], fld["chars"][7]
    if (
        False
        or (ch3, ch7) == ("D", "A")  # ..xDx.xAx..
        or (ch3, ch5) == ("D", "A")  # ..xDxAx.x..
        or (ch3, ch5) == ("C", "A")  # ..xCxAx.x..
        or (ch5, ch7) == ("D", "B")  # ..x.xDxBx..
        or (ch5, ch7) == ("D", "A")  # ..x.xDxAx..
    ):
        return float("inf")  # mutual blocking

    mfs = fld["spent"]
    d = {"A": 0, "B": 0, "C": 0, "D": 0}
    for pos, ch in enumerate(fld["chars"]):
        if ch == ".":
            continue
        if pos in ROOMS[ch]:
            if fld["guests"][ch]:
                mfs += (2 + 2) * (10 ** (ord(ch) - ord("A")))  # out and in
        else:
            roomnr = ord(ch) - ord("A")
            mfs += DISTANCE[pos][11 + roomnr + d[ch]] * (10 ** roomnr)
        d[ch] += 4
    return mfs


def selected_paths(fld):
    for path in SHORTEST_PATHS:
        frm, to_ = path[0], path[-1]
        if fld["chars"][frm] == "." or fld["chars"][to_] != ".":
            continue  # path must lead amphipod to empty space
        char = fld["chars"][frm]
        roomset = ROOMS[char]
        guests = fld["guests"][char]
        if frm in roomset and not guests:
            continue  # no need to leave
        if to_ in roomset:
            if guests:
                continue  # can't go home if guests still there
            if any(fld["chars"][k] == "." for k in roomset if k > to_):
                continue  # empty slots below to_
        else:
            if to_ not in range(11):
                continue  # forbidden room
        if any(fld["chars"][k] != "." for k in path[1:-1]):
            continue  # obstacles
        yield path


def try_paths(fld):
    global min_spent
    global seen_before

    for path in selected_paths(fld):
        new_fld = travel(fld, path)
        if new_fld["ident"] in seen_before:
            continue
        seen_before.add(new_fld["ident"])
        if new_fld["chars"][:11] == "...........":
            min_spent = min(min_spent, new_fld["spent"])
            continue
        if min_future_spent(new_fld) >= min_spent:
            continue  # don't recurse further
        try_paths(new_fld)


# construct playing field
data = open("input_23.txt").read()
lines = data.replace(" ", "").replace("#", "").splitlines()
chars = "..........." + lines[2] + "DCBA" + "DBAC" + lines[3]
fld = {
    "ident": "00000_" + chars,
    "spent": 0,
    "chars": chars,
    "guests": {
        ch: set(k for k in ROOMS[ch] if chars[k] not in {ch, "."})
        for ch in "ABCD"
    },
}

# solve the puzzle
seen_before = {fld["ident"]}
min_spent = float("inf")
try_paths(fld)
print(min_spent)