aoc2021-23.py

import functools
import sys

AMPHIPODS = "".join(c for c in sys.stdin.read() if c in "ABCD\n").strip().split("\n")
AMPHIPODS = [AMPHIPODS, [AMPHIPODS[0], "DCBA", "DBAC", AMPHIPODS[1]]]

PARTS = [
    tuple([tuple(["."] * 11)] + [tuple(row[i] for row in amphipods) for i in range(4)])
    for amphipods in AMPHIPODS
]


FORBIDDEN_BUFFER_POSITIONS = set(2 + i * 2 for i in range(4))
ROOMS = {chr(ord("A") + i): i for i in range(4)}
COSTS = {block: 10 ** room for block, room in ROOMS.items()}


def iter_possible_moves(state):
    # move out blocks from rooms to buffer space
    for room_idx, room in enumerate(state[1:]):
        expected = chr(ord("A") + room_idx)

        if not (set(room) - set([expected, "."])):
            # no blocks to move from this room
            continue

        # what are the possible positions in the buf space where we could put the blocks?
        room_pos = 2 + room_idx * 2
        left = right = room_pos
        while left >= 0 and state[0][left] == ".":
            left -= 1
        while right < len(state[0]) and state[0][right] == ".":
            right += 1

        possible_positions = set(range(left + 1, right))

        # never put blocks directly above the rooms
        possible_positions -= FORBIDDEN_BUFFER_POSITIONS

        # any spaces available?
        if not possible_positions:
            continue

        for depth, block in enumerate(room):
            if block == ".":
                continue

            step_cost = COSTS[block]

            for pos in possible_positions:
                steps = (depth + 1) + abs(room_pos - pos)
                yield (room_idx + 1, depth), (0, pos), step_cost * steps

            # only the first block from top can be moved
            break

    # put blocks back to rooms
    for pos, block in enumerate(state[0]):
        if block == ".":
            continue

        room_idx = ROOMS[block]

        if set(state[room_idx + 1]) - set([block, "."]):
            # there are conflicting blocks in the room, they need to be removed first
            continue

        room_pos = 2 + room_idx * 2

        if pos < room_pos:
            left, right = pos, room_pos
        else:
            right, left = pos, room_pos

        if any(state[0][i] != "." for i in range(left, right + 1) if i != pos):
            # block in the way
            continue

        # move possible!
        depth = max(i for i, c in enumerate(state[room_idx + 1]) if c == ".")
        step_cost = COSTS[block]
        steps = (depth + 1) + abs(room_pos - pos)
        yield (0, pos), (room_idx + 1, depth), step_cost * steps


def is_solved(state):
    return all(c == "." for c in state[0]) and all(
        set(room) == set([chr(ord("A") + room_idx)])
        for room_idx, room in enumerate(state[1:])
    )


def apply_move(state, src, dst):
    src_room, src_pos = src
    dst_room, dst_pos = dst
    # there must be a cleaner way to do this!
    state = [list(row) for row in state]
    state[dst_room][dst_pos], state[src_room][src_pos] = state[src_room][src_pos], "."
    return tuple(tuple(row) for row in state)


def draw(state):
    print("".join(state[0]))
    for row in range(len(state[1])):
        print("  " + " ".join(room[row] for room in state[1:]))
    print()


@functools.lru_cache(None)
def solve(state):
    if is_solved(state):
        return 0
    best_cost = float("Inf")
    for src, dst, move_cost in iter_possible_moves(state):
        if move_cost >= best_cost:
            continue
        new_state = apply_move(state, src, dst)
        cost = solve(new_state) + move_cost
        if best_cost > cost:
            best_cost = cost

    return best_cost


for part, state in enumerate(PARTS, 1):
    print(f"Solving part {part}:")
    draw(state)
    print(f"Min energy: {solve(state)}")
    print()