131 lines
3.7 KiB
Python
Executable file
131 lines
3.7 KiB
Python
Executable file
#!/usr/bin/env python
|
|
import heapq
|
|
import sys
|
|
from collections import defaultdict, deque
|
|
from dataclasses import dataclass
|
|
from functools import lru_cache
|
|
from math import inf
|
|
from typing import DefaultDict, Deque, Dict, FrozenSet, Iterator, List, Tuple, Union
|
|
|
|
RawGrid = List[str]
|
|
GraphInfo = List[Tuple[str, int]]
|
|
Graph = Dict[str, GraphInfo]
|
|
|
|
|
|
@dataclass(eq=True, frozen=True) # Hash-able
|
|
class Position:
|
|
x: int
|
|
y: int
|
|
|
|
|
|
def neighbours(grid: RawGrid, pos: Position) -> Iterator[Position]:
|
|
for dx, dy in ((1, 0), (-1, 0), (0, 1), (0, -1)):
|
|
new_pos = Position(pos.x + dx, pos.y + dy)
|
|
if not (0 <= new_pos.x < len(grid) and 0 <= new_pos.y < len(grid[0])):
|
|
continue
|
|
if grid[new_pos.x][new_pos.y] == "#":
|
|
continue
|
|
yield new_pos
|
|
|
|
|
|
def find_adjacent(grid: RawGrid, pos: Position) -> GraphInfo:
|
|
queue: Deque[Tuple[Position, int]] = deque()
|
|
visited = {pos}
|
|
adjacent: GraphInfo = []
|
|
|
|
for n in neighbours(grid, pos):
|
|
queue.append((n, 1)) # Distance is 1
|
|
|
|
while queue:
|
|
n, d = queue.popleft()
|
|
if n in visited:
|
|
continue
|
|
visited |= {n}
|
|
cell = grid[n.x][n.y]
|
|
|
|
if cell not in "#.1234": # We don't care about those
|
|
adjacent.append((cell, d))
|
|
continue # Do not go through doors and keys
|
|
|
|
for neighbour in neighbours(grid, n):
|
|
queue.append((neighbour, d + 1))
|
|
|
|
return adjacent
|
|
|
|
|
|
def build_graph(grid: RawGrid) -> Graph:
|
|
graph = {}
|
|
|
|
for x, row in enumerate(grid):
|
|
for y, cell in enumerate(row):
|
|
if cell not in "#.":
|
|
graph[cell] = find_adjacent(grid, Position(x, y))
|
|
|
|
return graph
|
|
|
|
|
|
def solve(G: Graph, start: str) -> int:
|
|
@lru_cache(2**20)
|
|
def reachable_keys(src: str, found: FrozenSet[str]) -> GraphInfo:
|
|
queue = []
|
|
distance: DefaultDict[str, Union[float, int]] = defaultdict(lambda: inf)
|
|
reachable: GraphInfo = []
|
|
|
|
for neighbor, weight in G[src]:
|
|
queue.append((weight, neighbor)) # Weight first for heap comparisons
|
|
|
|
heapq.heapify(queue)
|
|
|
|
while queue:
|
|
dist, node = heapq.heappop(queue)
|
|
|
|
# Do key, add it to reachable if not found previously
|
|
if node.islower() and node not in found:
|
|
reachable.append((node, dist))
|
|
continue
|
|
|
|
# Do door, if not opened by a key that was found in the search
|
|
if node.lower() not in found:
|
|
continue
|
|
|
|
# If not a key and not a closed door
|
|
for neighbor, weight in G[node]:
|
|
new_dist = dist + weight
|
|
if new_dist < distance[neighbor]:
|
|
distance[neighbor] = new_dist
|
|
heapq.heappush(queue, (new_dist, neighbor))
|
|
|
|
return reachable
|
|
|
|
@lru_cache(2**20)
|
|
def min_steps(
|
|
sources: str, keys_to_find: int, found: FrozenSet[str] = frozenset()
|
|
) -> Union[float, int]:
|
|
if keys_to_find == 0:
|
|
return 0
|
|
|
|
best = inf
|
|
|
|
for src in sources:
|
|
for key, dist in reachable_keys(src, found):
|
|
new_keys = found | {key}
|
|
new_sources = sources.replace(src, key)
|
|
new_dist = dist + min_steps(new_sources, keys_to_find - 1, new_keys)
|
|
|
|
if new_dist < best:
|
|
best = new_dist
|
|
|
|
return best
|
|
|
|
total_keys = sum(node.islower() for node in G)
|
|
return int(min_steps(start, total_keys)) # Throw if we kept the infinite float
|
|
|
|
|
|
def main() -> None:
|
|
G = build_graph(list(line.strip() for line in sys.stdin.readlines()))
|
|
print(solve(G, "1234"))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|