aoc/2022/14/solution.py

#!/usr/bin/python3
import sys

# part 1

def get_bounds(tiles):
    min_x = sys.maxsize
    min_y = sys.maxsize
    max_x = 0
    max_y = 0
    for x in tiles:
        for y in tiles[x]:
            if x < min_x: min_x = x
            if x > max_x: max_x = x
            if y < min_y: min_y = y
            if y > max_y: max_y = y
    return (min_x, min_y, max_x, max_y)

def print_grid(tiles, bounds):
    min_x, min_y, max_x, max_y = bounds
    for y in range(min_y, max_y + 1):
        for x in range(min_x, max_x + 1):
            if x in tiles and y in tiles[x]:
                print(tiles[x][y], end='')
            else:
                print(' ', end='')
        print()

def try_add(tiles, x, y, tile):
    if x not in tiles:
        tiles[x] = {y: tile}
    else:
        tiles[x][y] = tile

def collides_with(tiles, sand, pos):
    if px in tiles and py in tiles[px]:
        return True

def comes_to_rest(tiles, origin, bounds):
    min_x, min_y, max_x, max_y = bounds
    sx, sy = origin

    # lol
    if sx not in tiles:
        return False

    while sy <= max_y:
        if not (sx in tiles and sy + 1 in tiles[sx]):
            sy += 1
            continue
        elif not (sx - 1 in tiles and sy + 1 in tiles[sx - 1]):
            sx -= 1
            sy += 1
        elif not (sx + 1 in tiles and sy + 1 in tiles[sx + 1]):
            sx += 1
            sy += 1
        else:
            try_add(tiles, sx, sy, 'o')
            return True

    return False

def part1():
    tiles = {}
    for line in sys.stdin:
        new_rocks = [[int(pos) for pos in coord.split(',')] for coord in line.strip().split(' -> ')]
        print(new_rocks)
        for i, (rx1, ry1) in enumerate(new_rocks[1:]):
            rx2, ry2 = new_rocks[i]
            if rx1 - rx2 != 0:
                for x in range(min(rx1, rx2), max(rx1, rx2) + 1):
                    try_add(tiles, x, ry1, '#')
            elif ry1 - ry2 != 0:
                for y in range(min(ry1, ry2), max(ry1, ry2) + 1):
                    try_add(tiles, rx1, y, '#')
    print(tiles)

    bounds = get_bounds(tiles)
    rest = 0
    while comes_to_rest(tiles, (500, 0), bounds):
        rest += 1

    print_grid(tiles, bounds)
    print(f'{rest} units come to rest')

# part 2

def part2():
    tiles = {}
    for line in sys.stdin:
        new_rocks = [[int(pos) for pos in coord.split(',')] for coord in line.strip().split(' -> ')]
        print(new_rocks)
        for i, (rx1, ry1) in enumerate(new_rocks[1:]):
            rx2, ry2 = new_rocks[i]
            if rx1 - rx2 != 0:
                for x in range(min(rx1, rx2), max(rx1, rx2) + 1):
                    try_add(tiles, x, ry1, '#')
            elif ry1 - ry2 != 0:
                for y in range(min(ry1, ry2), max(ry1, ry2) + 1):
                    try_add(tiles, rx1, y, '#')
    print(tiles)

    min_x, min_y, max_x, max_y = get_bounds(tiles)
    least_x = 500 - (2 * (max_y - min_y))
    most_x = 500 + 2 * (max_y - min_y)
    for x in range(least_x, most_x + 1):
        try_add(tiles, x, max_y + 2, '#')
    bounds = get_bounds(tiles)

    print_grid(tiles, (min_x, min_y, max_x, max_y))
    rest = 0
    while not (500 in tiles and 0 in tiles[500]):
        comes_to_rest(tiles, (500, 0), bounds)
        rest += 1

    print()
    print_grid(tiles, (min_x, min_y, max_x, max_y))
    print(f'{rest} units come to rest')

if sys.argv[1] in '1':
    part1()
else:
    part2()