#!/usr/bin/env python import dataclasses import heapq import itertools import sys from collections.abc import Iterator from typing import NamedTuple NUM_FLOORS = 4 @dataclasses.dataclass(frozen=True) class Microchip: element: str @dataclasses.dataclass(frozen=True) class Generator: element: str Item = Microchip | Generator @dataclasses.dataclass(frozen=True, order=True) class State: class Floor(NamedTuple): chip: int generator: int elevator: int items: tuple[Floor, ...] def __post_init__(self) -> None: assert self.items == tuple(sorted(self.items)) # Sanity check def solve(input: str) -> int: def parse_item(input: str) -> Item: _, element, item_type = input.split() if item_type == "microchip": return Microchip(element.removesuffix("-compatible")) elif item_type == "generator": return Generator(element) assert False # Sanity check def parse_floor(input: str) -> list[Item]: # Simplify parsing, and remove `The Xth floor contains` input = input.removesuffix(".").replace(", and ", ", ").replace(" and ", ", ") input = " ".join(input.split()[4:]) if input == "nothing relevant": return [] return [parse_item(it) for it in input.split(", ")] def parse(input: str) -> dict[Item, int]: return { it: i for i, line in enumerate(input.splitlines()) for it in parse_floor(line) } def to_state(elevator: int, floors: dict[Item, int]) -> State: elements = {it.element for it in floors} return State( elevator=elevator, items=tuple( sorted( State.Floor(floors[Microchip(elem)], floors[Generator(elem)]) for elem in elements ) ), ) def from_state(state: State) -> tuple[int, dict[Item, int]]: floors: dict[Item, int] = {} for i, (chip, generator) in enumerate(state.items): floors[Microchip(str(i))] = chip floors[Generator(str(i))] = generator return state.elevator, floors def solve(state: State) -> int: def items_at( item_type: type[Item], floor: int, items: dict[Item, int], ) -> set[str]: return { it.element for it, it_floor in items.items() if it_floor == floor and isinstance(it, item_type) } def neighbours(state: State) -> Iterator[State]: elevator, all_items = from_state(state) chips = items_at(Microchip, elevator, all_items) gens = items_at(Generator, elevator, all_items) for dest_floor in (elevator - 1, elevator + 1): # Don't move the elevator out of bounds if dest_floor < 0 or dest_floor >= NUM_FLOORS: continue dest_chips = items_at(Microchip, dest_floor, all_items) dest_gens = items_at(Generator, dest_floor, all_items) unmatched_chips = dest_chips - dest_gens if unmatched_chips: assert not dest_gens # Sanity check single_items: list[Item] = [] # Chips for chip in chips: # can move to floors with no generator, or a matching generator if dest_gens and chip not in dest_gens: continue single_items.append(Microchip(chip)) # Generators for gen in gens: # can move to floors without unmatched chips or only their chip if unmatched_chips - {gen}: continue # ... but only if they're not currently protecting their chip if gen in chips and gens - {gen}: continue single_items.append(Generator(gen)) double_items: list[tuple[Item, Item]] = [] # Two chips for chip1, chip2 in itertools.combinations(chips, 2): # Can move to floors with no generator, or both matching generators if dest_gens and not dest_gens.issuperset({chip1, chip2}): continue double_items.append((Microchip(chip1), Microchip(chip2))) # Two generators for gen1, gen2 in itertools.combinations(gens, 2): # Can move to floors with unmatched chips, if they match them... if unmatched_chips - {gen1, gen2}: continue # ... but only if they're not currently protecting their chip if (gen1 in chips or gen2 in chips) and gens - {gen1, gen2}: continue double_items.append((Generator(gen1), Generator(gen2))) # Matching generator and chip for match in chips & gens: # Can move to floors with no unmatched chips if not unmatched_chips: double_items.append((Microchip(match), Generator(match))) for item in single_items: new_items = all_items | {item: dest_floor} assert new_items.keys() == all_items.keys() # Sanity check yield to_state(dest_floor, new_items) for item1, item2 in double_items: new_items = all_items | {item1: dest_floor, item2: dest_floor} assert new_items.keys() == all_items.keys() # Sanity check yield to_state(dest_floor, new_items) def dijkstra(start: State, end: State) -> int: # Priority queue of (distance, point) queue = [(0, start)] seen: set[State] = set() while len(queue) > 0: dist, p = heapq.heappop(queue) if p == end: return dist # We must have seen p with a smaller distance before if p in seen: continue # First time encountering p, must be the smallest distance to it seen.add(p) # Add all neighbours to be visited for n in neighbours(p): heapq.heappush(queue, (dist + 1, n)) assert False # Sanity check # On the end state, we want all items pairs on the top floor # The elevator must be on the top floor as well to get the last item up top = NUM_FLOORS - 1 end = State(top, tuple(State.Floor(top, top) for _ in state.items)) return dijkstra(state, end) floors = parse(input) for elem in ("elerium", "dilithium"): floors[Microchip(elem)] = 0 floors[Generator(elem)] = 0 state = to_state(0, floors) return solve(state) def main() -> None: input = sys.stdin.read() print(solve(input)) if __name__ == "__main__": main()