aoc2024day15.py

# === Parsing ===

obstacles = set()
walls = set()
initial_position = None
with open("input.txt", "r") as f:
    for y, line in enumerate(iter(f.readline, "\n")):
        for x, char in enumerate(line.strip()):
            if char == "@":
                initial_position = (x, y)
            elif char == "O":
                obstacles.add((x, y))
            elif char == "#":
                walls.add((x, y))

    moves = "".join(line.strip() for line in f)

assert initial_position is not None

print(moves[-10:])


# === Part 1 ===
def find_line_end(pos, dir, obstacles):
    px, py = pos
    dx, dy = dir
    while (px, py) in obstacles:
        px += dx
        py += dy
    return (px, py)


def move(pos, dir, obstacles, walls):
    px, py = pos
    dx, dy = dir
    nx, ny = px + dx, py + dy
    if (nx, ny) in walls:
        return (px, py), obstacles
    elif (nx, ny) in obstacles:
        line_end = find_line_end((nx, ny), dir, obstacles)
        if line_end in walls:
            return (px, py), obstacles
        else:
            obstacles.add(line_end)
            obstacles.remove((nx, ny))
            return (nx, ny), obstacles
    else:
        return (nx, ny), obstacles


DIRECTIONS = {
    ">": (1, 0),
    "v": (0, 1),
    "<": (-1, 0),
    "^": (0, -1),
}


pos = initial_position
for movement in moves:
    pos, obstacles = move(pos, DIRECTIONS[movement], obstacles, walls)

total_gps = 0
for x, y in obstacles:
    total_gps += 100 * y + x
print(total_gps)


# === Part 2 ===
from itertools import chain

obstacles = set()
walls = set()
initial_position = None
with open("input.txt", "r") as f:
    for y, line in enumerate(iter(f.readline, "\n")):
        for x, char in enumerate(line.strip()):
            if char == "@":
                initial_position = (2 * x, y)
            elif char == "O":
                obstacles.add((2 * x, y))
            elif char == "#":
                walls.add((2 * x, y))
                walls.add((2 * x + 1, y))

    moves = "".join(line.strip() for line in f)

assert initial_position is not None


def find_wide_line(pos, dx, obstacles):
    px, py = pos
    line = set()
    while (px, py) in obstacles:
        line.add((px, py))
        px += 2 * dx
    # Move one to the right if we were going left because
    # we skipped an empty position:
    # #....[][]@.. <
    # #..^ ^ ^     positions checked
    # #...^        return this one instead
    return (px if dx == 1 else px + 1, py), line


def expand_pushable_group(pos, dy, obstacles, walls):
    layer = {pos}
    group = {pos}
    while layer:
        # Check if this layer is hitting any walls.
        wall_positions = set(
            chain.from_iterable(((x, y + dy), (x + 1, y + dy)) for x, y in layer)
        )
        if wall_positions & walls:
            return False, None

        # Check the next layer.
        possible_next_layer_positions = set(
            chain.from_iterable(
                ((x - 1, y + dy), (x, y + dy), (x + 1, y + dy)) for x, y in layer
            )
        )
        new_layer = possible_next_layer_positions & obstacles
        group |= new_layer
        layer = new_layer

    return True, group


def move(pos, dir, obstacles, walls):
    px, py = pos
    dx, dy = dir
    nx, ny = px + dx, py + dy
    if (nx, ny) in walls:
        return (px, py), obstacles
    # Are we moving left/right?
    elif dx != 0 and ((nx, ny) in obstacles or (nx - 1, ny) in obstacles):
        obstacle_corner = ({(nx, ny), (nx - 1, ny)} & obstacles).pop()
        # Try to push obstacles left/right
        final_pos, line = find_wide_line(obstacle_corner, dx, obstacles)
        if final_pos in walls:
            return (px, py), obstacles
        else:
            pushed_line = {(x + dx, y) for x, y in line}
            obstacles = (obstacles - line) | pushed_line
            return (nx, ny), obstacles
    elif dy != 0 and ((nx, ny) in obstacles or (nx - 1, ny) in obstacles):
        # Try to push obstacles up/down
        obstacle_corner = ({(nx, ny), (nx - 1, ny)} & obstacles).pop()
        can_push, group = expand_pushable_group(obstacle_corner, dy, obstacles, walls)
        if not can_push:
            return (px, py), obstacles
        else:
            pushed_group = {(x, y + dy) for x, y in group}
            obstacles = (obstacles - group) | pushed_group
            return (nx, ny), obstacles
    else:
        return (nx, ny), obstacles


pos = initial_position
for movement in moves:
    pos, obstacles = move(pos, DIRECTIONS[movement], obstacles, walls)

total_gps = 0
for x, y in obstacles:
    total_gps += 100 * y + x
print(total_gps)