aoc-2022/14.py

lines = [line.strip() for line in open("14.input")]

rocks = set()

for line in lines:
    coords = line.split(" -> ")
    for i in range(0, len(coords) - 1):
        start = list(map(int, coords[i].split(",")))
        end = list(map(int, coords[i+1].split(",")))
        if start[0] == end[0]:
            for y in range(min(start[1], end[1]), max(start[1], end[1]) + 1):
                rocks.add((start[0], y))
        else:
            for x in range(min(start[0], end[0]), max(start[0], end[0]) + 1):
                rocks.add((x, start[1]))

def draw_grid(rocks, sands):
    obst = rocks.union(sands)
    min_x = min(o[0] for o in sands)
    max_x = max(o[0] for o in sands)
    max_y = max(o[1] for o in obst)
    for y in range(0, max_y+1):
        for x in range(min_x, max_x):
            if (x, y) in rocks:
                print("#", end="")
            elif (x, y) in sands:
                print("o", end="")
            else:
                print(".", end="")
        print("")

def step(rocks, sands, bottom):
    obst = rocks.union(sands)

    sand = (500, 0)

    if sand in sands:
        return sands

    while True:
        if sand[1] > bottom:
            return sands
        targets = [
           (sand[0], sand[1] + 1),
           (sand[0] - 1, sand[1] + 1),
           (sand[0] + 1, sand[1] + 1),
           ]
        for t in targets:
            if t not in obst:
                sand = t
                break
        else:
            break

    sands.add(sand)
    return sands

def part1():
    sands = set()
    bottom = max(r[1] for r in rocks)

    while True:
        l = len(sands)
        sands = step(rocks, sands, bottom)
        if l == len(sands):
            break

    print(len(sands))

part1()

def part2():
    sands = set()
    bottom = max(r[1] for r in rocks) + 2

    for x in range(-1000, 1000):
        rocks.add((x, bottom))

    while True:
        l = len(sands)
        sands = step(rocks, sands, bottom)
        if l == len(sands):
            break

    print(len(sands))

part2()