aoc-2021/20.py

from util import get_input

input = get_input("20.input")

enhancement = [c == '#' for c in input[0]]

img = set()
size = (len(input[2:][0]) - 1, len(input[2:]) - 1)

for y, line in enumerate(input[2:]):
    for x, c in enumerate(line):
        if c == '#':
            img.add((x, y))

def get_surrounding(i, size, stp, xin, yin):
    bnum = ""
    for y in [yin-1, yin, yin+1]:
        for x in [xin-1, xin, xin+1]:
            if (x, y) in i:
                bnum += "1"
            else:
                bnum += "0"
    res = int(bnum, 2)
    return res


def step(i, size, stp):
    newimg = set()
    for y in range(-stp - 1, size[1] + stp + 2):
        for x in range(-stp - 1, size[0] + stp + 2):
            idx = get_surrounding(i, size, stp, x, y)
            if enhancement[idx]:
                newimg.add((x, y))
    return newimg

def print_img(i, size, stp):
    for y in range(-stp - 1, size[1] + stp + 2):
        for x in range(-stp - 1, size[0] + stp + 2):
            print('#' if (x, y) in i else '.', end='')
        print("")

for i in range(50):
    #print_img(img, size, 50)
    #print("-" * 20)
    if i % 2 == 1:
        for y in range(-i - 2, size[1] + i + 1 + 2):
            img.add((-i - 2, y))
            img.add((-i - 1, y))
            img.add((size[0] + i + 1, y))
            img.add((size[0] + i + 2, y))
        for x in range(-i - 2, size[0] + i + 1 + 2):
            img.add((x, -i - 2))
            img.add((x, -i - 1))
            img.add((x, size[1] + i + 1))
            img.add((x, size[1] + i + 2))
    img = step(img, size, i)
    if i == 1:
        print("Part 1:", len(img))


#print_img(img, size, 50)

print("Part 2:", len(img))