aoc2024day10.py

# === Parsing ===
from collections import defaultdict

map = defaultdict(lambda: -1)
with open("input.txt", "r") as f:
    for y, line in enumerate(f):
        for x, char in enumerate(line.strip()):
            map[x, y] = int(char)
    MAX_X = x
    MAX_Y = y


print(map)


# === Part 1 ===
DIRECTIONS = [
    (0, 1),
    (1, 0),
    (0, -1),
    (-1, 0),
]


def score_head(map, pos):
    to_explore = [pos]
    seen = {pos}
    score = 0
    while to_explore:
        x, y = to_explore.pop()
        value = map[x, y]
        if value == 9:
            score += 1
            continue

        for dx, dy in DIRECTIONS:
            nx, ny = x + dx, y + dy
            if (nx, ny) in seen:
                continue
            if map[nx, ny] == value + 1:
                to_explore.append((nx, ny))
                seen.add((nx, ny))

    return score


from itertools import product

total = 0
for x, y in product(range(MAX_X + 1), range(MAX_Y + 1)):
    if map[x, y] == 0:
        total += score_head(map, (x, y))

print(total)


# === Part 2 ===
DIRECTIONS = [
    (0, 1),
    (1, 0),
    (0, -1),
    (-1, 0),
]


def rate_head(map, pos):
    to_explore = [pos]
    rating = 0
    while to_explore:
        x, y = to_explore.pop()
        value = map[x, y]
        if value == 9:
            rating += 1
            continue

        for dx, dy in DIRECTIONS:
            nx, ny = x + dx, y + dy
            if map[nx, ny] == value + 1:
                to_explore.append((nx, ny))

    return rating


from itertools import product

total = 0
for x, y in product(range(MAX_X + 1), range(MAX_Y + 1)):
    if map[x, y] == 0:
        total += rate_head(map, (x, y))

print(total)