problem10.py

import sys
from itertools import product

def first_pos(start, grid, m, n):
    x, y = start
    if x > 0 and grid[x-1][y] in ("|", "F", "7"):
        return (x-1, y), "N"
    if x < m-1 and grid[x+1][y] in ("|", "J", "L"):
        return (x+1, y), "S"
    if y > 0 and grid[x][y-1] in ("-", "F", "L"):
        return (x, y-1), "W"
    if y < n-1 and grid[x][y+1] in ("-", "7", "J"):
        return (x, y+1), "E"

def next_pos(pos, grid, m, n, prev_dir):
    x, y = pos
    match grid[x][y], prev_dir:
        case "F", "N":
            return (x, y+1), "E"
        case "F", "W":
            return (x+1, y), "S"
        case "|", "N":
            return (x-1, y), "N"
        case "|", "S":
            return (x+1, y), "S"
        case "-", "W":
            return (x, y-1), "W"
        case "-", "E":
            return (x, y+1), "E"
        case "L", "S":
            return (x, y+1), "E"
        case "L", "W":
            return (x-1, y), "N"
        case "7", "E":
            return (x+1, y), "S"
        case "7", "N":
            return (x, y-1), "W"
        case "J", "E":
            return (x-1, y), "N"
        case "J", "S":
            return (x, y-1), "W"

def follow(start, grid, m, n):
    loop = [start]
    pos, direction = first_pos(start, grid, m, n)
    loop.append(pos)
    while True:
        pos, direction = next_pos(pos, grid, m, n, direction)
        if pos == start:
            break
        else:
            loop.append(pos)
    return loop

def fill(start, marked, m, n, loop):
    x, y = start
    neighbors = [(x + 1, y), (x -1, y), (x, y-1), (x, y+1)]
    for neigh in neighbors:
        if neigh not in loop:
            newx, newy = neigh
            if 0 <= newx < m and 0 <= newy < n and neigh not in marked:
                marked.add(neigh)
                fill(neigh, marked, m, n, loop)

sys.setrecursionlimit(10000)
def inside(pos, grid, m, n, loop):
    x, y = pos
    count = 0
    while x >=0:
        x -= 1
        if (x, y) in loop:
            if grid[x][y] in ("-", "L", "F", "7", "J", "S"):
                count += 1
    return count % 2

with open(sys.argv[1]) as fh:
    grid = [line.rstrip() for line in fh]
    m, n = len(grid), len(grid[0])
    for i, r in enumerate(grid):
        if (j := r.find("S")) != -1:
            break
    start = (i, j)
    loop = follow(start, grid, m, n)
    x1, y1 = loop[1]
    x2, y2 = loop[-1]
    print((x1, y1), (x2, y2))
    breakpoint()
    loop
    loop = set(loop)
    buckets = []
    everything = set()
    while True:
        for i, j in product(range(m), range(n)):
            if (i, j) not in everything and (i, j) not in loop:
                break
        else:
            break
        marked = set([(i, j)])
        fill((i, j), marked, m, n, loop)
        buckets.append(marked)
        everything = everything.union(marked)
    size = 0
    for b in buckets:
        pos = b.pop()
        if inside(pos, grid, m, n, loop):
            print(b)
            size += len(b) + 1
    print(size)