day23_part1.py (15478B)
1 #!/usr/bin/env python 2 """Advent of Code 2021, day 23 (part 1): Amphipod 3 Play a shrimp shuffling game""" 4 5 # I think this is another problem that can be solved through graph search. The 6 # tricky bit (IMO) is that I can't just enumerate all of the possible states 7 # ahead of time and then plug them into SciPy or NetworkX and have them solve 8 # the problem. No, I think it'll be "easiest" to solve the graph and figure out 9 # states as we go. The next hard part is figuring out a _hashable_ 10 # representation of the state of the board. The example started off like this: 11 # 12 # ############# 13 # #...........# 14 # ###B#C#B#D### 15 # #A#D#C#A# 16 # ######### 17 # 18 # Another intermediate state looked like this: 19 # 20 # ############# 21 # #.....D.D.A.# 22 # ###.#B#C#.### 23 # #A#B#C#.# 24 # ######### 25 # 26 # There are 19 different spots through which an amphipod can pass (labeled 0 - 27 # I below), although they can't "stop" in all of them: 28 # 29 # ############# 30 # #0123456789A# 31 # ###B#D#F#H### 32 # #C#E#G#I# 33 # ######### 34 # 35 # And additionally, there are 8 amphipods, 2 of each type, with the different 36 # types having restrictions on which rooms they can enter and costs associated 37 # with moving. So while there may be some real overhead in manipulating this 38 # scheme, we can represent the state of the game using an 8-tuple that 39 # specifies the location where each amphipod is currently standing. 40 # 41 # The next big implementation decision is how to handle move restrictions, and 42 # I see two approaches: amphipods can either move to their final standing spot 43 # in one "step" of the algorithm (at the cost of more complex pathfinding 44 # logic), or they can take one step at a time, but extra logic can handle the 45 # fact that amphipods are "frozen" in place in the hallway until they can 46 # complete their move into their final room and can't stop in front of doors. 47 # I'm going to try the first approach since it seems like a smaller search 48 # problem. Basically, an amphipod can move from their starting room to one spot 49 # in the hallway, and then move to their final room, so maybe Dijkstra's 50 # algorithm would suffice; otherwise there are so many "low energy" moves to 51 # explore before the type D amphipods would ever move that we'd need A*, which 52 # means I'd need to implement a whole heuristic algorithm. The path-finding 53 # seems worth it. 54 55 from typing import Tuple, List, Union, Dict, Deque, NewType 56 from heapq import heappush, heappop 57 from collections import deque 58 59 from utils import get_puzzle_input 60 61 EXAMPLE_INPUT = \ 62 """############# 63 #...........# 64 ###B#C#B#D### 65 #A#D#C#A# 66 ######### 67 """ 68 69 GameState = NewType( 70 'GameState', 71 Tuple[int, int, int, int, int, int, int, int] 72 ) 73 74 def parse_input(input_string: str) -> GameState: 75 """Given the puzzle input, return the initial state of the game""" 76 # This implementation assumes that all amphipods start in a room; it's 77 # unable to parse a game state with amphipods in the hallway 78 lines = input_string.rstrip('\n').split('\n') 79 state_list: List[Union[None, int]] = [None] * 8 80 for room_line in range(2): 81 amphipods = [2 * (ord(c) - ord('A')) \ 82 for c in lines[2 + room_line].replace('#', '').replace(' ', '')] 83 for room_num, amphipod in enumerate(amphipods): 84 if state_list[amphipod] is not None: 85 amphipod += 1 86 state_list[amphipod] = room_line + 2 * room_num + 11 87 return GameState(tuple(state_list)) 88 89 # Because of my stubborn decision to represent the game state in an awkward 90 # way, I need a bunch of little functions to help me actually make sense of the 91 # state of the game. 92 93 def valid_hall_locations() -> Tuple[int, ...]: 94 """Return all of the hall locations in which an amphipod can legally stop""" 95 return (0, 1, 3, 5, 7, 9, 10) 96 97 def is_in_hall(location: int) -> bool: 98 """Returns `True` iff a location is in the hall""" 99 return location < 11 100 101 def get_room(location: int) -> Union[None, int]: 102 """Return the room number of the given location, or `None` if it's the 103 hallway""" 104 if is_in_hall(location): 105 return None 106 return (location - 11) // 2 107 108 def get_spot(location: int) -> Union[None, int]: 109 """Return the room spot number of the given location (0 at the top and 1 110 below it), or `None` if it's the hallway""" 111 if is_in_hall(location): 112 return None 113 return (location - 11) % 2 114 115 def get_entrance(room: int) -> int: 116 """Return the location of the hallway spot immediately outside a room""" 117 return 2 + 2 * room 118 119 def get_location(room: int, spot: int) -> int: 120 """Given a room and spot number, find the location""" 121 return 11 + room * 2 + spot 122 123 def is_empty(location: int, game_state: GameState) -> bool: 124 """Just indicates whether a location is empty; nothing special""" 125 return location not in game_state 126 127 def get_type(amphipod: int) -> int: 128 """There are 8 amphipods, but four types; this is useful because the type 129 ID is also the target room number""" 130 return amphipod // 2 131 132 def get_amphipod(location: int, game_state: GameState) -> Union[None, int]: 133 """Return the amphipod ID of the one in the given location, and `None` if 134 there is none""" 135 if location not in game_state: 136 return None 137 return game_state.index(location) 138 139 def get_energy_per_step(amphipod: int) -> int: 140 """Different types of amphipod consume different amounts of energy per 141 step""" 142 return (1, 10, 100, 1000)[get_type(amphipod)] 143 144 def is_home(amphipod: int, game_state: GameState) -> bool: 145 """An amphipod is 'home' if it's in its final room and there are no 146 amphipods below it in the room that aren't also 'home'""" 147 room = get_room(game_state[amphipod]) 148 # Now *this* is spaghetti code 149 if room is None: 150 return False 151 # Let's take a second to appreciate that tests for equality in Python work 152 # across object types, so this doesn't return None 153 if get_type(amphipod) == room: 154 spot = get_spot(game_state[amphipod]) 155 assert spot is not None 156 if spot == 1: 157 return True 158 next_location = get_location(room, spot+1) 159 if is_empty(next_location, game_state): 160 # The amphipod in question is in the right room, but there's 161 # weirdly an empty spot below it? 162 return False 163 return is_home(game_state.index(next_location), game_state) 164 return False 165 166 def everybody_home(game_state: GameState) -> bool: 167 """Return `True` iff every amphipod is in a home spot""" 168 return all(is_home(a, game_state) for a in range(len(game_state))) 169 170 def is_blocked_in_room(amphipod: int, game_state: GameState) -> bool: 171 """An amphipod is 'blocked' if they are in a room and another amphipod is 172 above them in the room (they might not be able to move anywhere in the 173 hallway, this function doesn't address that)""" 174 room = get_room(game_state[amphipod]) 175 if room is None: 176 return False 177 spot = get_spot(game_state[amphipod]) 178 if spot == 0: 179 return False 180 return not is_empty(get_location(room, 0), game_state) 181 182 def steps_between(from_location: int, to_location: int) -> int: 183 """Count the steps between a hall location and room location; one end must 184 be in a room and the other in the hallway""" 185 hall_location, room_location = sorted([from_location, to_location]) 186 assert is_in_hall(hall_location) and not is_in_hall(room_location) 187 return abs(hall_location - get_entrance(get_room(room_location))) \ 188 + get_spot(room_location) + 1 189 190 def path_clear(from_location: int, to_location: int, 191 game_state: GameState) -> bool: 192 """Return `True` iff there are no amphipods blocking the way between one 193 hallway location and another""" 194 # We allow an amphipod to be in the _from_ location, but not in the _to_ 195 # location (nor in any intermediate location) 196 assert is_in_hall(from_location) and is_in_hall(to_location) 197 if from_location == to_location: 198 return False 199 if from_location < to_location: 200 test_range = range(from_location+1, to_location+1) 201 else: 202 test_range = range(to_location, from_location) 203 return not any(i in game_state for i in test_range) 204 205 def make_state(game_state: GameState, amphipod: int, 206 new_location: int) -> GameState: 207 """Return a new state where the selected amphipod is in the given 208 location""" 209 game_state_list = list(game_state) 210 game_state_list[amphipod] = new_location 211 return GameState(tuple(game_state_list)) 212 213 # End of various helper functions. Next we're going to implement a version of 214 # Dijkstra's algorithm (that doesn't require the enumeration of every possible 215 # game state--which are the nodes in our graph). 216 # 217 # I'm going to use Python's built-in heap helpers and implement a version of a 218 # priority queue to support the algorithm. Normally, one of the trickiest 219 # operations in a heap-based priority queue is "removing" or "updating" entries 220 # from the queue (usually this is achieved by marking them as 'removed' somehow 221 # but otherwise leaving them in the queue). However, since queue entries are 222 # only being updated when a new distance to a node is *lower* than the previous 223 # one, that means anything that would be otherwise be removed/updated will 224 # appear as "already visited" when it's popped off the queue. In other words, 225 # as long as we don't consider states that have already been visited, we don't 226 # need to worry about removing/updating anything in the heap. 227 # 228 # In addition to the heap, we'll use a dictionary to quickly look up the 229 # current estimate of the minimum distance to a given node, and a set to keep 230 # track of visited nodes. Everything is built into Python. 231 232 def list_reachable_states(game_state: GameState) -> List[Tuple[GameState, int]]: 233 """Return a list for every (legal) game state that can be reached from the 234 current one, as a tuple containing the state itself and the energy that 235 would be required to move there""" 236 # This function, and not the simple implementation of Dijkstra's algorithm, 237 # is the meat of this solution, since there's a lot of tricky logic to 238 # figure out which amphipods can move and where they can move to. This 239 # finds all of the nodes connected to the current node in the graph. 240 reachable_states = [] 241 for amphipod, location in enumerate(game_state): 242 amphipod_type = get_type(amphipod) 243 amphipod_locations = [] 244 if is_in_hall(location): 245 # The only valid locations for a hallway amphipod to go to are its 246 # home room, and then only if it's empty or only occupied by 247 # amphipods of the same type. 248 if path_clear(location, get_entrance(amphipod_type), game_state): 249 for spot in (1, 0): 250 spot_location = get_location(amphipod_type, spot) 251 spot_amphipod = get_amphipod(spot_location, game_state) 252 if spot_amphipod is None: 253 amphipod_locations = [spot_location] 254 break 255 if get_type(spot_amphipod) != amphipod_type: 256 # The wrong type of amphipod is in this room so our guy 257 # can't enter it 258 break 259 elif not is_home(amphipod, game_state) \ 260 and not is_blocked_in_room(amphipod, game_state): 261 # If the amphipod is in a room (that's not its home), then the only 262 # valid location is somewhere in the hallway. 263 current_room_entrance = get_entrance(get_room(location)) 264 amphipod_locations = [l for l in valid_hall_locations() \ 265 if path_clear(current_room_entrance, l, 266 game_state)] 267 # else: pass 268 for new_location in amphipod_locations: 269 reachable_states.append(( 270 make_state(game_state, amphipod, new_location), 271 get_energy_per_step(amphipod) \ 272 * steps_between(location, new_location) 273 )) 274 return reachable_states 275 276 def reconstruct_path(start_state: GameState, final_state: GameState, 277 best_parent: Dict[GameState, GameState]) -> Deque[GameState]: 278 """Return the step-by-step solution to the problem""" 279 # I didn't need to implement this to solve the puzzle, but I did to debug 280 # my solution. :( 281 current_state = final_state 282 # Honestly a deque is overkill for this but I didn't use any for AoC and I 283 # wanted to appendleft so bad 284 path = deque([current_state]) 285 while current_state != start_state: 286 current_state = best_parent[current_state] 287 path.appendleft(current_state) 288 return path 289 290 def find_lowest_energy_use(start_state: GameState) -> Tuple[GameState, int, 291 Deque[GameState]]: 292 """Use Dijkstra's algorithm to find the lowest energy use to get all the 293 amphipods from the starting game state to a completed game state""" 294 priority_queue: List[Tuple[int, GameState]] = [] 295 visited_states = set() 296 best_distance: Dict[GameState, int] = {} 297 best_parent = {} 298 299 (energy_to_current_state, current_state) = (0, start_state) 300 301 while not everybody_home(current_state): 302 # One of the ways this departs from a usual graph traversal is that 303 # there are multiple "goal nodes" (game states). Any state that has 304 # every amphipod in its home location will do, and there's no point in 305 # enumerating them if we terminate the search once we find one. 306 visited_states.add(current_state) 307 308 # For each amphipod that can move, find all the states that are 309 # reachable from the current one. Add newly discovered nodes, or better 310 # than previously seen energy usage, to the queue. 311 for reachable_state, energy_between in list_reachable_states(current_state): 312 energy_to_reachable_state = energy_to_current_state \ 313 + energy_between 314 visit_neighbor = reachable_state not in visited_states and \ 315 (reachable_state not in best_distance or \ 316 best_distance[reachable_state] > energy_to_reachable_state) 317 if visit_neighbor: 318 heappush(priority_queue, (energy_to_reachable_state, reachable_state)) 319 best_distance[reachable_state] = energy_to_reachable_state 320 best_parent[reachable_state] = current_state 321 322 # Select the next lowest-energy state to visit next. The heap will 323 # become full of longer paths to various states that we've already 324 # visited, but we can skip past them quickly. 325 (energy_to_current_state, current_state) = heappop(priority_queue) 326 while current_state in visited_states: 327 (energy_to_current_state, current_state) = heappop(priority_queue) 328 329 return (current_state, energy_to_current_state, 330 reconstruct_path(start_state, current_state, best_parent)) 331 332 def solve_puzzle(input_string: str) -> int: 333 """Return the numeric solution to the puzzle""" 334 return find_lowest_energy_use(parse_input(input_string))[1] 335 336 def main() -> None: 337 """Run when the file is called as a script""" 338 assert solve_puzzle(EXAMPLE_INPUT) == 12521 339 print("Lowest energy use:", 340 solve_puzzle(get_puzzle_input(23))) 341 342 if __name__ == "__main__": 343 main()