day23_part2.py (18389B)
1 #!/usr/bin/env python 2 """Advent of Code 2021, day 23 (part 2): Amphipod 3 Play a shrimp shuffling game""" 4 5 # This began as a refactoring of the code that I wrote for part 1. Even before 6 # implementing a heuristic function (i.e., when this was essentially running 7 # Dijkstra's algorithm with a bit of pruning), this version ran much faster 8 # (260 vs 686 ms for the example input, and 1.68 vs. 4.1 s for my puzzle 9 # input). 10 # 11 # After adding a proper heuristic function, solution times dropped to 9.81 ms 12 # for the example input and 57 ms for my puzzle input. It takes 493 ms to solve 13 # part 2 using my puzzle input. 14 15 from typing import Tuple, List, Union, Dict, Deque, NewType, cast 16 from heapq import heappush, heappop 17 from collections import deque 18 from math import inf 19 20 from day23_part1 import EXAMPLE_INPUT 21 from utils import get_puzzle_input 22 23 ROOM_SIZE = 4 24 25 GameState = NewType( 26 'GameState', 27 Tuple[int, int, int, int, int, int, int, int, 28 int, int, int, int, int, int, int, int] 29 ) 30 31 def parse_input(input_string: str) -> GameState: 32 """Given the puzzle input, return the initial state of the game""" 33 # This implementation assumes that all amphipods start in a room; it's 34 # unable to parse a game state with amphipods in the hallway 35 lines = input_string.rstrip('\n').split('\n') 36 lines.insert(3, ' #D#C#B#A#') 37 lines.insert(4, ' #D#B#A#C#') 38 state_list: List[Union[None, int]] = [None] * 16 39 for room_line in range(4): 40 amphipods = [4 * (ord(c) - ord('A')) \ 41 for c in lines[2 + room_line].replace('#', '').replace(' ', '')] 42 for room_num, amphipod in enumerate(amphipods): 43 while state_list[amphipod] is not None: 44 amphipod += 1 45 state_list[amphipod] = room_line + 4 * room_num + 11 46 return cast(GameState, tuple(state_list)) 47 48 # Because of my stubborn decision to represent the game state in an awkward 49 # way, I need a bunch of little functions to help me actually make sense of the 50 # state of the game. 51 52 def valid_hall_locations() -> Tuple[int, ...]: 53 """Return all of the hall locations in which an amphipod can legally stop""" 54 #return tuple(l for l in range(11) if l not in range(2, 9, 2)) 55 return (0, 1, 3, 5, 7, 9, 10) 56 57 def room_entrance_locations() -> Tuple[Union[int, None], ...]: 58 """Return the location in the hallway of the entrance to the rooms""" 59 #return (None,) + tuple(range(2, 9, 2)) 60 return (None, 2, 4, 6, 8) 61 62 def everybody_home(game_state: GameState) -> bool: 63 """Quickly(ish) determine if all the amphipods are in their home rooms""" 64 return all(a // ROOM_SIZE == (l - 11) // ROOM_SIZE \ 65 for a, l in enumerate(game_state)) 66 67 def location_to_room_spot(location: int) -> Tuple[int, int]: 68 """Convert an 'absolute location' to a room and spot number (the hallway is 69 room 0)""" 70 if location < 11: 71 return (0, location) 72 return divmod(location - (11 - ROOM_SIZE), ROOM_SIZE) 73 74 def room_spot_to_location(room: int, spot: int) -> int: 75 """Given a room and spot number, find the location""" 76 if room == 0: 77 return spot 78 return 11 + (room-1) * ROOM_SIZE + spot 79 80 def amphipod_type(amphipod: int) -> int: 81 """There are 4 * `ROOM_SIZE` amphipods, but four types; this is useful 82 because the type ID is also the target room number (which is why this is 83 one-based)""" 84 return amphipod // ROOM_SIZE + 1 85 86 def get_energy_per_step(amphipod_home: int) -> int: 87 """Different types of amphipod consume different amounts of energy per 88 step""" 89 return cast(int, (None, 1, 10, 100, 1000)[amphipod_home]) 90 91 def make_state(game_state: GameState, amphipod: int, 92 new_location: int) -> GameState: 93 """Return a new state where the selected amphipod is in the given 94 location""" 95 game_state_list = list(game_state) 96 game_state_list[amphipod] = new_location 97 return cast(GameState, tuple(game_state_list)) 98 99 # Below are some functions (and their helpers) that return lists of potential 100 # moves. A "move" here means a location/number of steps pair (represented as a 101 # tuple). All of these return lists even though only `room_to_hallway` has the 102 # ability to return more than one move. 103 104 def stuck_in_room(room: int, spot: int, location_empty: List[bool]) -> bool: 105 """Helper function to indicate whether a given spot is "stuck" in a room""" 106 locations = (room_spot_to_location(room, s) for s in range(spot-1, -1, -1)) 107 # note to self: any([]) returns false and all([]) returns True 108 return not all(location_empty[l] for l in locations) 109 110 def path_clear(from_hall_spot: int, to_hall_spot: int, 111 location_empty: List[bool]) -> bool: 112 """Indicates whether the pathway is clear between the originating spot and 113 the destination spot (the originating spot need not be clear)""" 114 if from_hall_spot < to_hall_spot: 115 test_range = range(from_hall_spot + 1, to_hall_spot + 1) 116 else: 117 test_range = range(to_hall_spot, from_hall_spot) 118 return all(location_empty[l] for l in test_range) 119 120 def free_room_spot(room: int, location_empty: List[bool]) -> Union[int, None]: 121 """In a room that's ready, return the best (lowest) spot""" 122 for spot in range(ROOM_SIZE-1, -1, -1): 123 if location_empty[room_spot_to_location(room, spot)]: 124 return spot 125 return None 126 127 def room_to_hallway(room: int, spot: int, 128 location_empty: List[bool]) -> List[Tuple[int, int]]: 129 """List all the loecal moves from the given room/spot into the hallway""" 130 if stuck_in_room(room, spot, location_empty): 131 return [] 132 from_entrance = room_entrance_locations()[room] 133 assert from_entrance is not None 134 # This is extremely hacky, but we'll loop through the list of hall 135 # locations once in each direction. 136 destination_locations = [] 137 for location in valid_hall_locations(): 138 if location > from_entrance: 139 if not location_empty[location]: 140 break 141 destination_locations.append(location) 142 for location in reversed(valid_hall_locations()): 143 if location < from_entrance: 144 if not location_empty[location]: 145 break 146 destination_locations.append(location) 147 return [(l, abs(from_entrance - l) + spot + 1) \ 148 for l in destination_locations] 149 150 def hallway_to_room(hall_spot: int, room: int, 151 location_empty: List[bool]) -> List[Tuple[int, int]]: 152 """Return a list containing a move from the given spot in the hallway into 153 the given room, or an empty list if there's no such move""" 154 to_entrance = room_entrance_locations()[room] 155 assert to_entrance is not None 156 if not path_clear(hall_spot, to_entrance, location_empty): 157 return [] 158 to_spot = free_room_spot(room, location_empty) 159 assert to_spot is not None 160 return [(room_spot_to_location(room, to_spot), 161 to_spot + abs(hall_spot - to_entrance) + 1)] 162 163 def room_to_room(from_room: int, spot: int, to_room: int, 164 location_empty: List[bool]) -> List[Tuple[int, int]]: 165 """Return a list containing a move from the given spot and room into 166 the target room, or an empty list if there's no such move""" 167 from_entrance = room_entrance_locations()[from_room] 168 to_entrance = room_entrance_locations()[to_room] 169 assert from_entrance is not None and to_entrance is not None 170 if stuck_in_room(from_room, spot, location_empty) \ 171 or not path_clear(from_entrance, to_entrance, location_empty): 172 return [] 173 to_spot = free_room_spot(to_room, location_empty) 174 assert to_spot is not None 175 return [(room_spot_to_location(to_room, to_spot), 176 spot + to_spot + abs(from_entrance - to_entrance) + 2)] 177 178 # Now I finally implement A* and its most important helpers 179 180 def expand_game_state(game_state: GameState) -> \ 181 List[Tuple[int, int, int, int, int]]: 182 """Expand a game state tuple into a list of tuples specifying the room, 183 spot, home_room, amphipod index, and location of each amphipod, 184 reverse-sorted by room/spot""" 185 return sorted([(*location_to_room_spot(l), amphipod_type(a), a, l) \ 186 for a, l in enumerate(game_state)], 187 reverse=True) 188 189 def list_neighbors(game_state: GameState) -> List[Tuple[GameState, int]]: 190 """Return a list for almost every (legal) game state that can be reached from the 191 current one, as a tuple containing the state itself and the energy that 192 would be required to move there (there is some pruning here: if there are 193 paths that send amphipods home, only those neighbors are returned""" 194 # Before we start, we have to find all the _cozy_ amphipods and _ready_ 195 # rooms. A cozy amphipod is in its home room and isn't blocking any 196 # amphipod that's not. A ready room is a room in which any amphipods that 197 # are present are cozy. 198 game_state_expanded = expand_game_state(game_state) 199 location_empty = [True] * (11 + ROOM_SIZE * 4) 200 amphipod_cozy = [False] * len(game_state) 201 room_ready = [False] + [True] * 4 # hallway is never ready 202 last_was_cozy = False 203 for room, spot, home_room, amphipod, location in game_state_expanded: 204 location_empty[location] = False 205 if room == home_room: 206 if spot == ROOM_SIZE - 1: 207 amphipod_cozy[amphipod] = True 208 last_was_cozy = True 209 #room_ready[room] = True 210 elif last_was_cozy: 211 amphipod_cozy[amphipod] = True 212 elif room > 0: 213 last_was_cozy = False 214 room_ready[room] = False 215 216 # Now we'll loop through all the amphipods again and figure out where each 217 # can move. We'll build: 218 # * A list of all neighbor nodes to the current game state 219 reachable_states = [] 220 # * A subset of neighbors that put amphipods into rooms 221 reachable_room_states = [] 222 # Once we've found a neighbor state that puts an amphipod into a room, we 223 # won't bother with any other type of route 224 route_to_room = False 225 for room, spot, home_room, amphipod, _ in game_state_expanded: 226 legal_moves = [] 227 # We'll consider three situations in turn: 228 # * If the amphipod is in its home room, we'll move it to the hallway 229 # only if there is an amphipod of the wrong type below it (i.e., if 230 # it's not cozy). 231 # * If the amphipod is in the hallway, we'll move it to its home room 232 # only if any amphipods in the room are of the correct type (i.e, if 233 # the room is ready). 234 # * If the amphipod is in another room, we'll first check if there's a 235 # path to its home room (using the same rules as above); if not, then 236 # we'll consider moves into the hallway. 237 if room == home_room: 238 if not route_to_room and not amphipod_cozy[amphipod]: 239 # We need to consider moves to the hallway to get this amphipod 240 # out of the way of the amphipods below it (which need to move) 241 legal_moves = room_to_hallway(room, spot, location_empty) 242 elif room == 0: 243 if room_ready[home_room]: 244 # Consider moves to the home room (from the hallway) 245 legal_moves = hallway_to_room(spot, home_room, location_empty) 246 route_to_room = bool(legal_moves) 247 else: 248 if room_ready[home_room]: 249 legal_moves = room_to_room(room, spot, home_room, 250 location_empty) 251 route_to_room = bool(legal_moves) 252 if not route_to_room and not legal_moves: 253 # No route directly to home room, consider moves to the hallway 254 legal_moves = room_to_hallway(room, spot, location_empty) 255 256 for new_location, num_steps in legal_moves: 257 reachable_states.append(( 258 make_state(game_state, amphipod, new_location), 259 get_energy_per_step(home_room) * num_steps 260 )) 261 if route_to_room: 262 reachable_room_states.append(reachable_states[-1]) 263 264 if reachable_room_states: 265 return reachable_room_states 266 return reachable_states 267 268 def heuristic_to_end(game_state: GameState) -> int: 269 """A heuristic that estimates the total energy used between a given game 270 state and an ending condition""" 271 # We'll loop through amphipods twice: first to find out which amphipods are 272 # cozy (and their locations, which we'll mark as unavailable). 273 game_state_expanded = expand_game_state(game_state) 274 cozy_amphipods = set() 275 available_locations = set(range(11, 11 + 4 * ROOM_SIZE)) 276 last_was_cozy = False 277 for room, spot, home_room, amphipod, location in game_state_expanded: 278 if room == home_room: 279 if spot == ROOM_SIZE - 1: 280 cozy_amphipods.add(amphipod) 281 available_locations.remove(location) 282 last_was_cozy = True 283 elif last_was_cozy: 284 cozy_amphipods.add(amphipod) 285 available_locations.remove(location) 286 elif room > 0: 287 last_was_cozy = False 288 289 # Then, for the second loop, we'll put each amphipod into an available 290 # spot. Note that we're ignoring whether any other amphipods are in the way 291 # unless they're cozy. 292 energy = 0 293 for room, spot, home_room, amphipod, location in game_state_expanded: 294 if amphipod not in cozy_amphipods: 295 # Choose an available spot from their home room 296 for to_spot in range(ROOM_SIZE): 297 to_location = room_spot_to_location(home_room, to_spot) 298 if to_location in available_locations: 299 available_locations.remove(to_location) 300 break 301 else: 302 raise RuntimeError('no spot for amphipod') 303 304 # Find the steps/energy to this spot 305 num_steps = to_spot + 1 306 if room == 0: 307 hall_spot = spot 308 else: 309 hall_spot = room_entrance_locations()[room] 310 assert hall_spot is not None 311 num_steps += spot + 1 312 if room == home_room: 313 # This is a non-cozy amphipod in its own room, so it'll need to 314 # leave the room and return 315 num_steps += 2 316 num_steps += abs(hall_spot - room_entrance_locations()[home_room]) 317 energy += num_steps * get_energy_per_step(home_room) 318 return energy 319 320 def reconstruct_path(came_from: Dict[GameState, GameState], 321 current_state: GameState) -> Deque[GameState]: 322 """Return the step-by-step solution to the problem""" 323 # I didn't need to implement this to solve the puzzle, but I did to debug 324 # my solution. :( 325 # Also, a deque is overkill for this but I didn't use any for AoC and I 326 # wanted to (efficiently) appendleft so bad 327 total_path = deque([current_state]) 328 while current_state in came_from: 329 current_state = came_from[current_state] 330 total_path.appendleft(current_state) 331 return total_path 332 333 def find_lowest_energy_use(start_state: GameState) -> Tuple[GameState, int, 334 Deque[GameState]]: 335 """Use the A* algorithm to find the lowest energy use to get all the 336 amphipods from the starting game state to a completed game state""" 337 # Influenced heavily by the priority queue implementation in Python's heapq 338 # documentation. Note that passing multiple types into List[] is 339 # technically incorrect, but we're using lists instead of types 340 # specifically to get pass-by-reference behavior, so I'm accepting the mypy 341 # errors as a cost of clearer documentation. 342 open_set: List[List[int, Union[Tuple[int], GameState]]] = [] 343 open_set_entries: Dict[GameState, List[int, GameState]] = {} 344 345 came_from: Dict[GameState, GameState] = {} 346 best_distance: Dict[GameState, Union[int, float]] = {start_state: 0} 347 best_cost: Dict[GameState, Union[int, float]] = { 348 start_state: 0 + heuristic_to_end(start_state) 349 } 350 351 current_state = start_state 352 while not everybody_home(current_state): 353 for neighbor, energy_between in list_neighbors(current_state): 354 energy_to_neighbor = best_distance[current_state] \ 355 + energy_between 356 # If the distance to the neighbor from the start through the 357 # current node is lower than what was previously estimated, or if 358 # we've never seen a path to this neighbor before, update the 359 # priority queue and other data sctructures. 360 if energy_to_neighbor < best_distance.get(neighbor, inf): 361 estimated_cost = energy_to_neighbor \ 362 + heuristic_to_end(neighbor) 363 came_from[neighbor] = current_state 364 best_distance[neighbor] = energy_to_neighbor 365 best_cost[neighbor] = estimated_cost 366 # Updating the priority queue is a bit more complicated 367 new_pq_entry = [estimated_cost, neighbor] 368 if neighbor in open_set_entries: 369 old_pq_entry = open_set_entries.pop(neighbor) 370 # Placeholder for "empty" priority queue entries 371 old_pq_entry[-1] = (0,) 372 open_set_entries[neighbor] = new_pq_entry 373 heappush(open_set, new_pq_entry) 374 375 # Select the next lowest-energy state to visit next. 376 while open_set: 377 _, current_state = heappop(open_set) 378 if current_state != (0,): 379 del open_set_entries[current_state] 380 break 381 else: 382 raise KeyError('pop from empty priority queue') 383 384 return (current_state, int(best_distance[current_state]), 385 reconstruct_path(came_from, current_state)) 386 387 def solve_puzzle(input_string: str) -> int: 388 """Return the numeric solution to the puzzle""" 389 return find_lowest_energy_use(parse_input(input_string))[1] 390 391 def main() -> None: 392 """Run when the file is called as a script""" 393 assert solve_puzzle(EXAMPLE_INPUT) == 44169 394 print("Lowest energy use:", 395 solve_puzzle(get_puzzle_input(23))) 396 397 if __name__ == "__main__": 398 main()