aoc2024day12.py

# === Parsing ===

from collections import defaultdict

garden = defaultdict(str)
with open("input.txt", "r") as f:
    for y, line in enumerate(f):
        for x, plot in enumerate(line.strip()):
            garden[x, y] = plot
    WIDTH = x + 1
    HEIGHT = y + 1


# === Part 1 ===
from itertools import product


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


def get_region(garden, pos):
    """Given a position, finds the whole region of the same plot."""
    to_explore = [pos]
    region = set()
    while to_explore:
        x, y = to_explore.pop()
        plot = garden[x, y]
        region.add((x, y))
        for dx, dy in DIRECTIONS:
            nx, ny = x + dx, y + dy
            if garden[nx, ny] == plot and (nx, ny) not in region:
                to_explore.append((nx, ny))

    return region


def compute_perimeter(region):
    """Computes the perimeter of a region."""
    return sum(
        (x + dx, y + dy) not in region for x, y in region for dx, dy in DIRECTIONS
    )


seen = set()
regions = defaultdict(list)
for x, y in product(range(WIDTH), range(HEIGHT)):
    if (x, y) in seen:
        continue

    region = get_region(garden, (x, y))
    seen.update(region)
    regions[garden[x, y]].append(region)

cost = 0
for region_list in regions.values():
    for region in region_list:
        area = len(region)
        perimeter = compute_perimeter(region)
        cost += area * perimeter

print(cost)


# === Part 2 ===
def count_sides(region):
    side_segments = {
        ((x, y), (dx, dy))
        for x, y in region
        for dx, dy in DIRECTIONS
        if (x + dx, y + dy) not in region
    }

    sides = 0
    while side_segments:
        # New side to try to extend.
        sides += 1
        (x, y), (dir_x, dir_y) = side_segments.pop()
        dx = 0 if dir_x else 1
        dy = 0 if dir_y else 1
        for sign in [-1, 1]:
            step = 1
            while True:
                nx, ny = x + step * sign * dx, y + step * sign * dy
                possible_segment = ((nx, ny), (dir_x, dir_y))
                if possible_segment in side_segments:
                    side_segments.remove(possible_segment)
                    step += 1
                else:
                    break

    return sides


cost = 0
for region_list in regions.values():
    for region in region_list:
        area = len(region)
        sides = count_sides(region)
        cost += area * sides

print(cost)