aoc2024day6.py

# === Parsing ===

obstacles = set()  # (x, y) positions of the obstacles.
initial_position = None
with open("input.txt", "r") as f:
    for y, line in enumerate(f):
        for x, char in enumerate(line):
            if char == "#":
                obstacles.add((x, y))
            elif char == "^":
                initial_position = (x, y)

    MAX_X = x
    MAX_Y = y

assert initial_position is not None

print(obstacles, initial_position)


# === Part 1 ===

from itertools import cycle

directions = cycle(
    [
        (0, -1),  # up
        (1, 0),  # right
        (0, 1),  # down
        (-1, 0),  # left
    ]
)
dx, dy = next(directions)

positions = set()
x, y = initial_position
while True:
    positions.add((x, y))
    nx, ny = x + dx, y + dy
    if nx < 0 or nx > MAX_X or ny < 0 or ny > MAX_Y:
        break
    if (nx, ny) in obstacles:
        dx, dy = next(directions)
    else:
        x, y = nx, ny

print(len(positions))


# === Part 2 ===


def has_loop(obstacles, initial_position):
    directions = cycle([(0, -1), (1, 0), (0, 1), (-1, 0)])
    dx, dy = next(directions)

    positions = set()
    x, y = initial_position
    while True:
        positions.add((x, y, dx, dy))
        nx, ny = x + dx, y + dy
        if nx < 0 or nx > MAX_X or ny < 0 or ny > MAX_Y:
            return False
        if (nx, ny) in obstacles:
            dx, dy = next(directions)
        elif (nx, ny, dx, dy) in positions:
            return True
        else:
            x, y = nx, ny


positions.remove(initial_position)
count = 0
for pos in positions:
    # Put a new temporary obstacle
    obstacles.add(pos)
    # Check if there's a loop:
    count += has_loop(obstacles, initial_position)
    # Remove the temporary obstacle
    obstacles.remove(pos)

print(count)