https://adventofcode.com/2023/day/16

day16.py

import itertools as it

L = [l.strip() for l in open(0)]
N, M = len(L), len(L[0])
starts: list[tuple[int, int]] = list(
    it.chain(*zip(*(((i, -1), (i, M)) for i in range(N))),
             *zip(*(((-1, j), (N, j)) for j in range(M)))))
ms = [(1, 0), (0, 1), (-1, 0), (0, -1)]


def dfs(x, y, dx, dy, E, V) -> None:
    if (x, y, dx, dy) in V: return
    V.add((x, y, dx, dy))

    while 0<=(x := x+dx)<N and 0<=(y := y+dy)<M:
        E[x][y] = 1
        if L[x][y]!='.': break

    if x<0 or x>=N or y<0 or y>=M: return

    c = L[x][y]
    if c=='-' and dx:
        dfs(x, y, 0, 1, E, V)
        dfs(x, y, 0, -1, E, V)
    elif c=='|' and dy:
        dfs(x, y, 1, 0, E, V)
        dfs(x, y, -1, 0, E, V)
    elif c=='/':
        dfs(x, y, -dy, -dx, E, V)
    elif c=='\\':
        dfs(x, y, dy, dx, E, V)
    else:
        dfs(x, y, dx, dy, E, V)


def energize(x, y, dx, dy) -> int:
    E = [[0]*M for _ in range(N)]
    dfs(x, y, dx, dy, E, set())
    return sum(sum(r) for r in E)


print(energize(0, -1, 0, 1))  # part 1
print(max(energize(*s, *ds) for s, ds in it.product(starts, ms)))  # part 2