advent/y2019/p20.py

import itertools
import sys

from toolkit import read_image, shortest_path


def move(pos):
    for adj in (pos + ori for ori in [1, -1, 1j, -1j]):
        if adj in grid:
            yield adj
        if adj in duals:
            yield from move(duals[adj])


def move_with_depth(state):
    pos, level = state
    for adj in (pos + ori for ori in [1, -1, 1j, -1j]):
        if adj in inner:
            yield from move_with_depth((duals[adj], level + 1))
        elif adj in outer and (level > 0) and (grid[adj] not in {'AA', 'ZZ'}):
            yield from move_with_depth((duals[adj], level - 1))
        elif adj in grid:
            yield adj, level


# base
text = sys.stdin.read()
grid = {pos: val for pos, val in read_image(text).items() if val not in ' #'}
for pos, val in grid.copy().items():
    if val.isupper():
        for im in [1, 1j]:
            A, B, C = [pos - im, pos, pos + im]
            if grid.get(C) == '.':
                grid[B] = grid.pop(A) + grid.pop(B)
            if grid.get(A) == '.':
                grid[B] = grid.pop(B) + grid.pop(C)
elements = {v: k for k, v in grid.items()}

# matrix
duals = {}
portals = {pos for pos, cell in grid.items() if cell.isupper()}
for pA, pB in itertools.combinations(portals, 2):
    if grid[pA] == grid[pB]:
        duals[pA] = pB
        duals[pB] = pA

# part1
start, end = elements['AA'], elements['ZZ']
path = shortest_path(start, end, move)
print(len(path) - 3)

# categorize portals
_, *xmid, _ = sorted({p.real for p in portals})
_, *ymid, _ = sorted({p.imag for p in portals})
inner = {pos for pos in portals if pos.real in xmid and pos.imag in ymid}
outer = set(portals) - inner

# part2
start, end = (elements['AA'], 0), (elements['ZZ'], 0)
path = shortest_path(start, end, move_with_depth)
print(len(path) - 3)