advent-of-code/2021/d23/ex1/ex1.py

201 lines
5.2 KiB
Python
Executable file

#!/usr/bin/env python
import enum
import functools
import sys
from typing import Iterator, List, NamedTuple, Optional, Tuple, cast
class Point(NamedTuple):
x: int
y: int
class Amphipod(enum.IntEnum):
A = 0
B = 1
C = 2
D = 3
class Direction(enum.IntEnum):
ALLEY = 0
ROOM = 1
# 7-length tuple, but easier for mypy in a variadic type
Alley = Tuple[Optional[Amphipod], ...]
# Actually variadic tuple, instead of a list, for memoization purposes
Room = Tuple[Amphipod, ...]
# 4-length tuple, but easier for mypy in a variadic type
Rooms = Tuple[Room, ...]
class Board(NamedTuple):
alley: Alley
rooms: Rooms
class Move(NamedTuple):
cost: int
new_board: Board
FUEL_COST = {
Amphipod.A: 1,
Amphipod.B: 10,
Amphipod.C: 100,
Amphipod.D: 1000,
}
DISTANCE = [
# From room 1
(2, 1, 1, 3, 5, 7, 8),
# From room 2
(4, 3, 1, 1, 3, 5, 6),
# From room 3
(6, 5, 3, 1, 1, 3, 4),
# From room 4
(8, 7, 5, 3, 1, 1, 2),
]
AMPHIPOD_FROM_STRING = {
"A": Amphipod.A,
"B": Amphipod.B,
"C": Amphipod.C,
"D": Amphipod.D,
}
ROOM_SIZE = 2
def solve(input: List[str]) -> int:
def parse() -> Board:
alley: Alley = (None,) * 7
rooms: Rooms = ()
for i in (3, 5, 7, 9):
room: Room = tuple(
AMPHIPOD_FROM_STRING[input[j][i]] for j in range(2, 3 + 1)
)
rooms = rooms + (room,)
return Board(alley, rooms)
def room_is_solved(board: Board, amphipod: Amphipod) -> bool:
room = board.rooms[amphipod]
return len(room) == ROOM_SIZE and all(a == amphipod for a in room)
def board_is_solved(board: Board) -> bool:
return all(room_is_solved(board, Amphipod(i)) for i in range(len(board.rooms)))
def move_cost(
board: Board, room: int, alley_spot: int, direction: Direction
) -> Optional[int]:
# Going left-to-right, or right-to-left
if room < (alley_spot - 1):
alley_start = room + 2
# Look at the end spot if we're going to the alley, not if we come from there
alley_end = alley_spot + (1 - direction)
else:
# Look at the first spot if we're going to the alley, not if we come from there
alley_start = alley_spot + direction
alley_end = room + 2
# Is there any obstacle in the way
if any(spot is not None for spot in board.alley[alley_start:alley_end]):
return None
amphipod = (
board.alley[alley_spot]
if direction == Direction.ROOM
else board.rooms[room][0]
)
assert amphipod is not None # Sanity check
return FUEL_COST[amphipod] * (
DISTANCE[room][alley_spot] + direction + ROOM_SIZE - len(board.rooms[room])
)
# Yes this returns a 0-or-1 length iterator, but it's practical for `moves`
def alley_moves_for(board: Board, i: int) -> Iterator[Move]:
# Return early if we're trying to move out of an empty spot
spot = board.alley[i]
if spot is None:
return
# Can't yet move to the target room if any amphipod is out of place there
if any(other != spot for other in board.rooms[spot]):
return
cost = move_cost(board, spot, i, Direction.ROOM)
# Can't move there yet, there's an obstacle in the way
if cost is None:
return
# Update the board state
alley, rooms = board
rooms = rooms[:spot] + ((spot,) + rooms[spot],) + rooms[spot + 1 :]
alley = alley[:i] + (None,) + alley[i + 1 :]
yield Move(cost, Board(alley, rooms))
def rooms_moves_for(board: Board, i: int) -> Iterator[Move]:
room = board.rooms[i]
# No need to move out of a solved room
if all(a == i for a in room):
return
for dest in range(len(board.alley)):
cost = move_cost(board, i, dest, Direction.ALLEY)
# Can't move there yet, there's an obstacle in the way
if cost is None:
continue
# Update the board state
alley, rooms = board
rooms = rooms[:i] + (room[1:],) + rooms[i + 1 :]
alley = alley[:dest] + (room[0],) + alley[dest + 1 :]
yield Move(cost, Board(alley, rooms))
def moves(board: Board) -> Iterator[Move]:
for i in range(len(board.alley)):
yield from alley_moves_for(board, i)
for i in range(len(board.rooms)):
yield from rooms_moves_for(board, i)
@functools.cache
def total_cost(board: Board) -> Optional[int]:
if board_is_solved(board):
return 0
best = None
for cost, new_board in moves(board):
if (end_cost := total_cost(new_board)) is None:
continue
cost += end_cost
if best is None or cost < best:
best = cost
return best
board = parse()
cost = total_cost(board)
assert cost is not None # Sanity check
return cost
def main() -> None:
input = [line.rstrip("\n") for line in sys.stdin.readlines()]
print(solve(input))
if __name__ == "__main__":
main()