advent-of-code/2023/d16/ex1/ex1.py

119 lines
3.2 KiB
Python
Executable file

#!/usr/bin/env python
import sys
from collections import deque
from collections.abc import Iterator
from enum import Enum, StrEnum
from typing import NamedTuple
class Point(NamedTuple):
x: int
y: int
class Direction(Enum):
NORTH = Point(-1, 0)
SOUTH = Point(1, 0)
EAST = Point(0, 1)
WEST = Point(0, -1)
def apply(self, pos: Point) -> Point:
dx, dy = self.value
return Point(pos.x + dx, pos.y + dy)
class Tile(StrEnum):
EMPTY = "."
MIRROR = "/"
ANTI_MIRROR = "\\"
VERTICAL_SPLITTER = "|"
HORIZONTAL_SPLITTER = "-"
def apply(self, dir: Direction) -> Iterator[Direction]:
match self:
case self.EMPTY:
yield dir
case self.MIRROR:
directions = {
Direction.EAST: Direction.NORTH,
Direction.NORTH: Direction.EAST,
Direction.WEST: Direction.SOUTH,
Direction.SOUTH: Direction.WEST,
}
yield directions[dir]
case self.ANTI_MIRROR:
directions = {
Direction.WEST: Direction.NORTH,
Direction.NORTH: Direction.WEST,
Direction.EAST: Direction.SOUTH,
Direction.SOUTH: Direction.EAST,
}
yield directions[dir]
case self.VERTICAL_SPLITTER:
if dir in (Direction.NORTH, Direction.SOUTH):
yield dir
else:
yield from (Direction.NORTH, Direction.SOUTH)
case self.HORIZONTAL_SPLITTER:
if dir in (Direction.EAST, Direction.WEST):
yield dir
else:
yield from (Direction.EAST, Direction.WEST)
class TileMap(NamedTuple):
tiles: dict[Point, Tile]
lines: int
rows: int
def in_bounds(self, pos: Point) -> bool:
x, y = pos
if x < 0 or x >= self.lines:
return False
if y < 0 or y >= self.rows:
return False
return True
def solve(input: list[str]) -> int:
def parse(input: list[str]) -> TileMap:
res: dict[Point, Tile] = {}
for x, line in enumerate(input):
for y, c in enumerate(line):
res[Point(x, y)] = Tile(c)
return TileMap(res, len(input), len(input[0]))
def compute_energy(map: TileMap) -> int:
queue: deque[tuple[Point, Direction]] = deque([(Point(0, 0), Direction.EAST)])
seen: set[tuple[Point, Direction]] = set()
energized: set[Point] = set()
while queue:
point, dir = queue.popleft()
if (point, dir) in seen:
continue
seen.add((point, dir))
energized.add(point)
for new_dir in map.tiles[point].apply(dir):
new_pos = new_dir.apply(point)
if not map.in_bounds(new_pos):
continue
queue.append((new_pos, new_dir))
return len(energized)
map = parse(input)
return compute_energy(map)
def main() -> None:
input = sys.stdin.read().splitlines()
print(solve(input))
if __name__ == "__main__":
main()