advent-of-code/2022/python/day14.py

import matrix
import shared
import itertools
import functools


SAND_START = [0, 500]


def find_bounds(lines):
    x = []
    y = []
    for line in lines:
        x.extend(matrix.split_x_out(line))
        y.extend(matrix.split_y_out(line))
    return shared.minmax(y), shared.minmax(x)


def part1(rows):
    def get_next(sand):
        next_y = sand[0] + 1
        dl = sand[1] - 1
        d = sand[1]
        dr = sand[1] + 1
        return next_y, (dl, d, dr)

    lines = [
        [
            list(reversed(list(map(int, point.split(",")))))
            for point in row.split(" -> ")
        ]
        for row in rows
    ]
    (minY, maxY), (minX, maxX) = find_bounds(lines)
    mx = matrix.matrix_of_size(maxX + 2, maxY + 2)
    walls = []
    for line in lines:
        for x in range(len(line) - 1):
            coords = matrix.coords_between_points(line[x], line[x + 1])
            walls.extend(coords)
    for y, x in walls:
        mx[y][x] = "#"
    mx[SAND_START[0]][SAND_START[1]] = "+"

    sands = []

    try:
        for sand_grains in range(900):  # produce Sand one unit at a time
            last_grain = []
            sand = SAND_START[:]
            # print("starting sand", sand_grains+1, sand)
            while True:
                last_grain.append(sand[:])
                next_y, (dl, d, dr) = get_next(sand)
                if mx[next_y][d] == 0:
                    # print(mx[next_y][d],"moving_down", end=' ')
                    # print("moving down")
                    sand[0] = next_y
                elif mx[next_y][dl] == 0:
                    # print(mx[next_y][dl], "moving left", end=' ')
                    # print("moving left")
                    sand[0] = next_y
                    sand[1] = dl
                elif mx[next_y][dr] == 0:
                    # print(mx[next_y][dr], "moving right", end=' ')
                    # print("moving right")
                    sand[0] = next_y
                    sand[1] = dr
                else:
                    mx[sand[0]][sand[1]] = "o"
                    # print(mx[next_y][dl], mx[next_y][d], mx[next_y][dr])
                    # print("STOPPING", end=' ')
                    break
            # matrix.view_matrix(mx, 0,minX-1, maxY+1,maxX+2)
    except IndexError:
        sands.append(last_grain)
        for y, x in sands[-1]:
            mx[y][x] = "~"
        # matrix.view_matrix(mx, 0,minX-1, maxY+1,maxX+2)
        print("done", sand_grains)


@shared.profile
def part2(rows):
    lines = [
        [
            list(reversed(list(map(int, point.split(",")))))
            for point in row.split(" -> ")
        ]
        for row in rows
    ]
    (minY, maxY), (minX, maxX) = find_bounds(lines)
    walls = {}
    _between_points = matrix.coords_between_points
    for line in lines:
        for x in range(len(line) - 1):
            coords = matrix.coords_between_points(line[x], line[x + 1])
            for c in coords:
                walls[c] = "#"
    for y, x in walls:
        walls[(y, x)] = "#"  # .add((y,x))
    floor = maxY + 2

    sand_grains = -1
    while True:
        sand_grains += 1
        sand = [0, 500]

        if tuple(sand) in walls:
            print("done", sand_grains)
            break

        while True:
            next_y = sand[0] + 1
            dl = sand[1] - 1
            dr = sand[1] + 1
            if next_y == floor:
                # hit the floor
                walls[tuple(sand)] = "#"  # Draw stop
                break

            elif (next_y, sand[1]) not in walls:
                sand[0] = next_y
            elif (next_y, dl) not in walls:
                sand[0] = next_y
                sand[1] = dl
            elif (next_y, dr) not in walls:
                sand[0] = next_y
                sand[1] = dr

            else:
                # Hit something flat
                walls[tuple(sand)] = "#"
                break


def main():
    rows = [row for row in shared.load_rows(14)]
    with shared.elapsed_timer() as elapsed:
        part1(rows)
        print("🕒", elapsed())
    with shared.elapsed_timer() as elapsed:
        part2(rows)
        print("🕒", elapsed())


if __name__ == "__main__":
    main()