1  from lib import *
 2  
 3  input = read_input(2021, 23)
 4  
 5  
 6  def generate_moves(rooms, hallway, n):
 7      def check_hallway(start, end):
 8          return all(hallway[i] is None or i == start for i in range(min(start, end), max(start, end) + 1))
 9  
10      def push_room(idx, elem):
11          return rooms[:idx] + (rooms[idx] + (elem,),) + rooms[idx + 1 :]
12  
13      def pop_room(idx):
14          return rooms[:idx] + (rooms[idx][:-1],) + rooms[idx + 1 :]
15  
16      def set_hallway(idx, elem):
17          return hallway[:idx] + (elem,) + hallway[idx + 1 :]
18  
19      for i, c in enumerate(hallway):
20          if c is None:
21              continue
22  
23          dst = "ABCD".index(c)
24  
25          if any(x != c for x in rooms[dst]):
26              continue
27  
28          if not check_hallway(i, 2 + 2 * dst):
29              continue
30  
31          dist = abs(2 + 2 * dst - i) + (n - len(rooms[dst]))
32  
33          yield dist * 10**dst, push_room(dst, c), set_hallway(i, None)
34  
35          return
36  
37      for i in range(4):
38          if all(x == "ABCD"[i] for x in rooms[i]):
39              continue
40  
41          c = rooms[i][-1]
42  
43          src = 2 + 2 * i
44  
45          dst = "ABCD".index(c)
46  
47          for j in [0, 1, 3, 5, 7, 9, 10]:
48              if not check_hallway(src, j):
49                  continue
50  
51              dist = (1 + n - len(rooms[i])) + abs(src - j)
52  
53              yield dist * 10**dst, pop_room(i), set_hallway(j, c)
54  
55  
56  def solve(part2):
57      lines = input.splitlines()[3:1:-1]
58  
59      if part2:
60          lines.insert(1, "  #D#B#A#C#")
61  
62          lines.insert(2, "  #D#C#B#A#")
63  
64      n = len(lines)
65  
66      queue = [(0, cnt := 0, tuple([*zip(*lines)][3:-1:2]), (None,) * 11)]
67  
68      visited = set()
69  
70      while queue:
71          energy, _, rooms, hallway = heappop(queue)
72  
73          if (rooms, hallway) in visited:
74              continue
75  
76          visited.add((rooms, hallway))
77  
78          if rooms == tuple((c,) * n for c in "ABCD"):
79              return energy
80  
81          for c, r, h in generate_moves(rooms, hallway, n):
82              cnt += 1
83  
84              heappush(queue, (energy + c, cnt, r, h))
85  
86  
87  print(solve(False))
88  print(solve(True))