adventofcode2023/day10/sol.py

import sys
from collections import defaultdict
map = [x.strip() for x in sys.stdin]

print(map)

Y = len(map)
X = len(map[0])
assert all(len(row) == X for row in map)

N,E,S,W = 1,2,4,8
unknown = 0b11111
grid = {}
start = None
for i, row in enumerate(map):
    for j,c in enumerate(row):
        d = 0
        if c == ".": d = 0
        elif c == "7": d = W|S
        elif c == "J": d = N|W
        elif c == "L": d = N|E
        elif c == "F": d = S|E
        elif c == "|": d = N|S
        elif c == "-": d = W|E
        elif c == "S":
            d = unknown
            start = (i,j)
        else:
            raise Exception("invalid char %s at (%s,%s)" % (c, i, j))
        grid[i,j] = d

#print(grid)

def find_loop(grid, start):
    front = [(0,start)]
    dist = {}
    count = defaultdict(int)
    count[start] -= 1
    def push(i,j,from_):
        if grid.get((i,j),0)&from_:
            front.append((cost+1, (i,j)))
    while front:
        front.sort()
        cost, n = front.pop(0)
        if n not in grid:
            continue
        count[n] += 1
        if n in dist:
            continue
        dist[n] = cost
        i, j = n
        d = grid[n]
        if d&N: push(i-1,j,S)
        if d&S: push(i+1,j,N)
        if d&W: push(i,j-1,E)
        if d&E: push(i,j+1,W)

    #print(count.values())
    if not all(c == 2 for c in count.values()):
        # not a loop
        return {}

    return dist

def enclosedby(grid, points):
    # upsample by 3x3
    import numpy
    fill = numpy.zeros((Y*3,X*3+1), dtype='uint8')
    #fill = [[0]*(X*3) for _ in range(Y*3)]
    for i,j in points:
        d = grid[i,j]
        if   d == N|S: t=[[0,1,0],[0,1,0],[0,1,0]]
        elif d == N|W: t=[[0,1,0],[1,1,0],[0,0,0]]
        elif d == N|E: t=[[0,1,0],[0,1,1],[0,0,0]]
        elif d == S|W: t=[[0,0,0],[1,1,0],[0,1,0]]
        elif d == S|E: t=[[0,0,0],[0,1,1],[0,1,0]]
        elif d == E|W: t=[[0,0,0],[1,1,1],[0,0,0]]
        else: raise Exception("invalid d=%d at (%s,%s)" % (d, i, j))
        for k in range(3):
            for l in range(3):
                fill[i*3+k, j*3+l] = t[k][l]

    def draw(fill):
        for row in fill:
            print("".join(".+O"[f] for f in row))

    draw(fill)

    # fill edges
    for i in range(Y*3):
        if fill[i,   0  ] == 0: fill[i,   0  ] = 2
        if fill[i, X*3-1] == 0: fill[i, X*3-1] = 2
    for j in range(X*3):
        if fill[  0  , j] == 0: fill[  0  , j] = 2
        if fill[Y*3-1, j] == 0: fill[Y*3-1, j] = 2

    # flood fill
    changed = True
    while changed:
        #draw(fill)
        changed = False
        for i in range(Y*3):
            f = (fill[i] == 2)
            f = numpy.roll(f, 1) | numpy.roll(f, -1)
            if i > 0:     f |= (fill[i-1] == 2)
            if i < Y*3-1: f |= (fill[i+1] == 2)
            f[-1] = False
            f &= (fill[i] == 0)
            if f.any():
                fill[i, f] = 2
                changed = True
    draw(fill)

    # downsample
    inside = []
    for i,j in grid:
        f = [fill[i*3+k, j*3+l]==0 for k in range(3) for l in range(3)]
        if all(f):
            inside.append((i,j))


    print("inside =", inside)
    return len(inside)

for D in [W|S, N|W, N|E, S|E, N|S, W|E]:
    grid[start] = D
    loop = find_loop(grid, start)
    if loop:
        print(D, loop.keys())
        print(max(loop.values()))
        print(enclosedby(grid, loop))