adventofcode2024/day12/sol.py

import itertools
def solve(file):
    grid = [l.strip() for l in file]
    ny = len(grid)
    nx = len(grid[0])
    def get(p):
        x,y = p
        if 0 <= x < nx and 0 <= y < ny:
            return grid[y][x]
        return '.'

    def flood(initial_point):
        c = get(initial_point)
        stack = [initial_point]
        count = {} # set of points connected to x,y. value is the number of connected neighboring points
        while stack:
            p = stack.pop()
            if p not in count:
                count[p] = 0
                x,y = p
                for n in (x+1,y), (x-1,y), (x,y+1), (x,y-1):
                    if get(n) == c:
                        stack.append(n)
                        count[p] += 1
        return count

    points = lambda: itertools.product(range(nx), range(ny))

    done = set() # set of points which have been counted as part of a shape
    total = 0
    total2 = 0
    for p in points():
        if p in done: continue
        shape = flood(p)
        done |= shape.keys()
        area = len(shape)
        perim = 0
        for p,n in shape.items():
            perim += 4-n
        #print(get(p), area, perim)
        total += area * perim
        total2 += area * discount(shape)

    print(total)
    print(total2)

def discount(shape):
    def sides(p):
        x,y = p
        s = 0
        if (x+1,y) in shape: s += 1
        if (x-1,y) in shape: s += 2
        if (x,y+1) in shape: s += 4
        if (x,y-1) in shape: s += 8
        return s

    # only count a side as free if the neighbor to the north or west
    # does not have the same side free
    perim = 0
    for p in shape:
        x,y = p
        free = sides(p) ^ 15
        if (x-1,y) in shape:
            free &= sides((x-1,y))
        if (x,y-1) in shape:
            free &= sides((x,y-1))
        perim += free.bit_count()
    return perim

solve(open("sample4.in"))
solve(open("input"))