adventofcode2022/day18/sol.py

import numpy

N = 23

def solve():
    lava = numpy.zeros((N,N,N), dtype=numpy.int32)
    for line in open("input"):
        x,y,z = map(int, line.strip().split(","))
        lava[x,y,z] = 1

    lava = numpy.pad(lava, (1,1), constant_values=0)

    # part 1
    area = count_touching(lava, lava^1)
    print(area)

    # part 2
    border = numpy.zeros((N,N,N), dtype=numpy.int32)
    border = numpy.pad(border, (1,1), constant_values=1)

    water = flood(lava, border)
    #print(water)

    area = count_touching(lava, water)
    print(area)


def count_touching(A, B):
    """counts the number of places a cell in A touches a cell in B
    (assumes a 1-cell padding around the edge which is not counted)"""
    n = 0
    for x in range(1,N+1):
        for y in range(1,N+1):
            for z in range(1,N+1):
                if A[x,y,z]:
                    n += B[x-1,y,z]
                    n += B[x+1,y,z]
                    n += B[x,y-1,z]
                    n += B[x,y+1,z]
                    n += B[x,y,z-1]
                    n += B[x,y,z+1]
    return n

def flood(lava, border):
    def roll(a, axis, amount):
        return numpy.roll(a, amount, axis=axis)

    a = border
    while True:
        x = roll(a, 0, -1)
        x |= roll(a, 0, +1)
        x |= roll(a, 1, -1)
        x |= roll(a, 1, +1)
        x |= roll(a, 2, -1)
        x |= roll(a, 2, +1)
        x &= ~lava
        b = a | x
        if (a == b).all():
            return a
        a = b

solve()