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
    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
    next = {}
    def set(i,j):
        if 0 <= i < Y*3 and 0 <= j < X*3:
            if fill[i][j] == 0:
                next[i,j] = 2
    while True:
        #draw(fill)
        next.clear()
        for i in range(Y*3):
            for j in range(X*3):
                if fill[i][j] == 2:
                    set(i-1,j)
                    set(i+1,j)
                    set(i,j-1)
                    set(i,j+1)
        if next:
            for (i,j),c in next.items():
                fill[i][j] = c
            continue
        break
    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))