2019: d18: ex2: add solution

2019/d18/ex2/ex2.py

@@ -0,0 +1,131 @@
+#!/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()