aoc/2023/16/solution.py

#!/usr/bin/python3
import sys
from dataclasses import dataclass

@dataclass
class Pos:
    y: int
    x: int

names = ('up', 'right', 'down', 'left')
emittance = {
    '.': ([], [], [], []),
    '|': ([], [0, 2], [], [0, 2]),
    '-': ([1, 3], [], [1, 3], []),
    '\\': ([3], [2], [1], [0]),
    '/': ([1], [0], [3], [2]),
}

def step(pos, direction):
    return (Pos(pos.y - 1, pos.x), Pos(pos.y, pos.x + 1), Pos(pos.y + 1, pos.x), Pos(pos.y, pos.x - 1))[direction]

def bounded(pos, bounds):
    return pos.y >= 0 and pos.x >= 0 and pos.y < bounds.y and pos.x < bounds.x

def collide(grid, energized, bounds, ray, direction):
    if (ray.y, ray.x, direction) in energized: return
    while bounded(ray, bounds):
        energized.add((ray.y, ray.x, direction))
        emit = emittance[grid[ray.y][ray.x]][direction]
        for r in emit:
            collide(grid, energized, bounds, step(ray, r), r)
        if len(emit) > 0:
            return
        ray = step(ray, direction)

def energy_grid_str(energized, bounds):
    seen = set()
    for y, x, direction in energized:
        seen.add((y, x))
    return '\n'.join([''.join(row) for row in [['#' if (y, x) in seen else ' ' for x in range(bounds.x)] for y in range(bounds.y)]])

def energy_grid_count(energized, bounds):
    seen = set()
    for y, x, direction in energized:
        seen.add((y, x))
    return len(seen)

def part1():
    print('part 1')
    grid = [tuple(line) for line in map(str.strip, sys.stdin)]
    energized = set()
    bounds = Pos(len(grid), len(grid[0]))
    ray = Pos(0, 0)
    collide(grid, energized, bounds, ray, 1)
    print(energy_grid_str(energized, bounds))
    print(f'{energy_grid_count(energized, bounds)} tiles energized')

def part2():
    print('part 2')
    grid = [tuple(line) for line in map(str.strip, sys.stdin)]
    bounds = Pos(len(grid), len(grid[0]))
    max_energized = 0
    for y in range(0, bounds.y):
        energized = set()
        collide(grid, energized, bounds, Pos(y, 0), 1)
        result = energy_grid_count(energized, bounds)
        max_energized = max(max_energized, result)
        print(f'{result} tiles energized')
    for x in range(0, bounds.x):
        energized = set()
        collide(grid, energized, bounds, Pos(0, x), 2)
        result = energy_grid_count(energized, bounds)
        max_energized = max(max_energized, result)
        print(f'{result} tiles energized')
    for y in range(0, bounds.y):
        energized = set()
        collide(grid, energized, bounds, Pos(y, bounds.x - 1), 3)
        result = energy_grid_count(energized, bounds)
        max_energized = max(max_energized, result)
        print(f'{result} tiles energized')
    for x in range(0, bounds.x):
        energized = set()
        collide(grid, energized, bounds, Pos(bounds.y - 1, x), 0)
        result = energy_grid_count(energized, bounds)
        max_energized = max(max_energized, result)
        print(f'{result} tiles energized')
    print(f'maximum is {max_energized}')

if sys.argv[1] in '1':
    part1()
else:
    part2()