2019: d20: ex1: add solution

This commit is contained in:
Bruno BELANYI 2024-12-26 21:39:32 -05:00
parent 944a4cc401
commit e8e90e64ed

106
2019/d20/ex1/ex1.py Executable file
View file

@ -0,0 +1,106 @@
#!/usr/bin/env python
import heapq
import sys
from collections import defaultdict
from typing import Iterator, NamedTuple
class Point(NamedTuple):
x: int
y: int
def neighbours(self) -> Iterator["Point"]:
for dx, dy in (
(-1, 0),
(1, 0),
(0, -1),
(0, 1),
):
yield Point(self.x + dx, self.y + dy)
Graph = dict[Point, set[Point]]
def solve(input: str) -> int:
def post_process_gates(
letters: dict[Point, str], paths: set[Point]
) -> dict[str, set[Point]]:
res: dict[str, set[Point]] = defaultdict(set)
for p1, first in letters.items():
for dx, dy in ((0, 1), (1, 0)):
p2 = Point(p1.x + dx, p1.y + dy)
if p2 not in letters:
continue
gate = first + letters[p2]
p0 = Point(p1.x - dx, p1.y - dy)
p3 = Point(p2.x + dx, p2.y + dy)
res[gate] |= {p0, p3} & paths
return res
def to_graph(paths: set[Point], gates: dict[str, set[Point]]) -> Graph:
res: dict[Point, set[Point]] = defaultdict(set)
for p in paths:
res[p] |= set(p.neighbours()) & paths
for gate, points in gates.items():
if len(points) == 1:
assert gate in ("AA", "ZZ") # Sanity check
continue
for p in points:
res[p] |= points
res[p].remove(p)
return res
def parse(input: list[str]) -> tuple[Graph, Point, Point]:
letters: dict[Point, str] = {}
paths: set[Point] = set()
for x, line in enumerate(input):
for y, c in enumerate(line):
if c == "#" or c == " ":
continue
p = Point(x, y)
if c == ".":
paths.add(p)
continue
letters[p] = c
gates = post_process_gates(letters, paths)
graph = to_graph(paths, post_process_gates(letters, paths))
return graph, next(iter(gates["AA"])), next(iter(gates["ZZ"]))
def djikstra(start: Point, end: Point, graph: Graph) -> int:
# Priority queue of (distance, point)
queue = [(0, start)]
seen: set[Point] = 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 graph[p]:
heapq.heappush(queue, (dist + 1, n))
assert False # Sanity check
graph, start, end = parse(input.splitlines())
return djikstra(start, end, graph)
def main() -> None:
input = sys.stdin.read()
print(solve(input))
if __name__ == "__main__":
main()