Advent of code 2021 day 23 compact

def process1(s, part2=False):
  lines = s.strip('\n').split('\n')
  letters = [[line[3], line[5], line[7], line[9]] for line in lines[2:4]]
  if part2:
    letters[1:1] = [list('DCBA'), list('DBAC')]  # Insert two middle rows.
  nrows = len(letters)
  # State consists of the 7 eligible hallway locations and the room locations.
  start_state = tuple(['.'] * 7 + list(itertools.chain(*letters)))
  end_state = tuple(['.'] * 7 + list('ABCD') * nrows)
  cost_for_id = {'A': 1, 'B': 10, 'C': 100, 'D': 1000}

  @functools.lru_cache(maxsize=None)
  def compute_cost(state):
    if state == end_state:
      return 0
    min_cost = 10**8

    # Explore moving an item from position i0 to i1, and update min cost.
    def consider(i0, i1, distance):
      nonlocal min_cost
      id = state[i0]  # id in 'ABCD'
      state2 = list(state)
      state2[i0] = '.'
      state2[i1] = id  # (previously '.')
      candidate_cost = distance * cost_for_id[id] + compute_cost(tuple(state2))
      min_cost = min(min_cost, candidate_cost)

    # Move from the hallway position `i0` to a room `j`.
    for i0 in range(7):
      id = state[i0]
      if id not in 'ABCD':
        continue
      j = ord(id) - 65
      row = max((i for i in range(nrows) if state[7 + i * 4 + j] == '.'),
                default=None)
      if (row is None or
          not all(state[7 + i * 4 + j] == id for i in range(row + 1, nrows))):
        continue
      i1 = 7 + row * 4 + j
      move_right = i0 <= j + 1
      if all(state[k] == '.' for k in (
          range(i0 + 1, 2 + j) if move_right else range(2 + j, i0))):
        distance = ((j - i0 + 2) * 2 - (i0 == 0) + row if move_right else
                    (i0 - j - 1) * 2 - (i0 == 6) + row)
        consider(i0, i1, distance)

    # Move from a room `j` to the hallway position `i1`.
    for j in range(4):
      row = min((i for i in range(nrows) if state[7 + i * 4 + j] != '.'),
                default=None)
      if row is None:
        continue
      i0 = 7 + row * 4 + j
      id = state[i0]
      if (id in ['.', chr(65 + j)] and not any(
          state[7 + i * 4 + j] != id for i in range(row + 1, nrows))):
        continue
      for i1 in range(2 + j, 7):  # Move right in hallway.
        if all(state[k] == '.' for k in range(2 + j, i1 + 1)):
          consider(i0, i1, (i1 - j - 1) * 2 - (i1 == 6) + row)
      for i1 in range(0, 2 + j):  # Move left in hallway
        if all(state[k] == '.' for k in range(i1, 2 + j)):
          consider(i0, i1, (j - i1 + 2) * 2 - (i1 == 0) + row)

    return min_cost

  return compute_cost(start_state)