|
| 1 | +# -*- coding: utf-8 -*- |
| 2 | + |
| 3 | +from collections import defaultdict |
| 4 | +from heapq import heappush, heappop |
| 5 | + |
| 6 | + |
| 7 | +def get_diagram_from_input(data): |
| 8 | + return data.split("\n") |
| 9 | + |
| 10 | + |
| 11 | +def read_file(filename): |
| 12 | + with open(filename) as file: |
| 13 | + data = file.read() |
| 14 | + return get_diagram_from_input(data) |
| 15 | + |
| 16 | + |
| 17 | +def get_state_from_diagram(diagram): |
| 18 | + state = defaultdict(set) |
| 19 | + for y, line in enumerate(diagram): |
| 20 | + for x, char in enumerate(line): |
| 21 | + if char in "ABCD": |
| 22 | + state[char].add((x, y)) |
| 23 | + return tuple(pos for char in "ABCD" for pos in sorted(state[char])) |
| 24 | + |
| 25 | + |
| 26 | +def perfrom_amphipod_in_room(x, y, positions, rooms, char, state, num_lines, hallway, new): |
| 27 | + if (x, y - 1) in positions: |
| 28 | + return False |
| 29 | + |
| 30 | + if rooms[char] == x: |
| 31 | + if all( |
| 32 | + (x, y_) in positions |
| 33 | + and char == "ABCD"[state.index((x, y_)) // num_lines] |
| 34 | + for y_ in range(y + 1, num_lines + 2) |
| 35 | + ): |
| 36 | + return False |
| 37 | + |
| 38 | + left_hallway = hallway[hallway.index(x - 1):: -1] |
| 39 | + right_hallway = hallway[hallway.index(x + 1):] |
| 40 | + for hallway_ in (left_hallway, right_hallway): |
| 41 | + for x_ in hallway_: |
| 42 | + if (x_, 1) in positions: |
| 43 | + break |
| 44 | + new.append((x_, 1)) |
| 45 | + return new |
| 46 | + |
| 47 | + |
| 48 | +def perfrom_amphipod_in_hallway(num_lines, rooms, char, positions, state, x, new): |
| 49 | + skip = False |
| 50 | + for y_ in range(num_lines + 1, 1, -1): |
| 51 | + if (rooms[char], y_) in positions: |
| 52 | + c_ = "ABCD"[state.index((rooms[char], y_)) // num_lines] |
| 53 | + if c_ != char: |
| 54 | + skip = True |
| 55 | + break |
| 56 | + else: |
| 57 | + break |
| 58 | + if skip: |
| 59 | + return False |
| 60 | + if x < rooms[char]: |
| 61 | + hallway_ = range(x + 1, rooms[char] + 1) |
| 62 | + else: |
| 63 | + hallway_ = range(x - 1, rooms[char] - 1, - 1) |
| 64 | + for x_ in hallway_: |
| 65 | + if (x_, 1) in positions: |
| 66 | + break |
| 67 | + else: |
| 68 | + assert rooms[char] == x_ |
| 69 | + new.append((x_, y_)) |
| 70 | + return new |
| 71 | + |
| 72 | + |
| 73 | +def get_least_energy_organize_amphipods(diagram, num_lines): |
| 74 | + state = get_state_from_diagram(diagram) |
| 75 | + rooms = {"A": 3, "B": 5, "C": 7, "D": 9} |
| 76 | + hallway = [1, 2, 4, 6, 8, 10, 11] |
| 77 | + seen = {state: 0} |
| 78 | + visit = list([(0, state)]) |
| 79 | + while visit: |
| 80 | + _, state = heappop(visit) |
| 81 | + cost = seen[state] |
| 82 | + if state == tuple((i, j) for i in range(3, 10, 2) for j in range(2, num_lines + 2)): |
| 83 | + return cost |
| 84 | + positions = set(state) |
| 85 | + for i in range(len(state)): |
| 86 | + new = [] |
| 87 | + char = "ABCD"[i // num_lines] |
| 88 | + x, y = state[i] |
| 89 | + if y != 1: |
| 90 | + response = perfrom_amphipod_in_room(x, y, positions, rooms, char, state, num_lines, hallway, new) |
| 91 | + if response is False: |
| 92 | + continue |
| 93 | + else: |
| 94 | + new = response |
| 95 | + else: |
| 96 | + response = perfrom_amphipod_in_hallway(num_lines, rooms, char, positions, state, x, new) |
| 97 | + if response is False: |
| 98 | + continue |
| 99 | + else: |
| 100 | + new = response |
| 101 | + |
| 102 | + L = i // num_lines * num_lines |
| 103 | + R = L + num_lines |
| 104 | + left = state[:L] |
| 105 | + right = state[R:] |
| 106 | + mid = state[L:R] |
| 107 | + for x_, y_ in new: |
| 108 | + my_amphipods = tuple(sorted(tuple(xy for xy in mid if xy != (x, y)) + ((x_, y_),))) |
| 109 | + state_ = left + my_amphipods + right |
| 110 | + move_cost = pow(10, "ABCD".index(char)) |
| 111 | + new_cost = cost + (abs(x - x_) + abs(y - y_)) * move_cost |
| 112 | + if state_ not in seen or seen[state_] > new_cost: |
| 113 | + seen[state_] = new_cost |
| 114 | + heappush(visit, (new_cost, state_)) |
| 115 | + |
| 116 | + |
| 117 | +if __name__ == "__main__": |
| 118 | + diagram = read_file("input.txt") |
| 119 | + |
| 120 | + first_solution = get_least_energy_organize_amphipods(diagram, num_lines=2) |
| 121 | + print(first_solution) |
| 122 | + |
| 123 | + new_diagram = diagram[:3] + [" #D#C#B#A#", " #D#B#A#C#"] + diagram[3:] |
| 124 | + second_solution = get_least_energy_organize_amphipods(new_diagram, num_lines=4) |
| 125 | + print(second_solution) |
0 commit comments