Spaces:
Running
Running
file = "input.txt" | |
def read_input(file): | |
with open(file, 'r') as f: | |
return [line.strip() for line in f] | |
def get_neighbors(x, y, max_x, max_y): | |
neighbors = [] | |
if x > 0: | |
neighbors.append((x - 1, y)) | |
if x < max_x - 1: | |
neighbors.append((x + 1, y)) | |
if y > 0: | |
neighbors.append((x, y - 1)) | |
if y < max_y - 1: | |
neighbors.append((x, y + 1)) | |
return neighbors | |
def flood_fill(grid, x, y, visited): | |
plant_type = grid[x][y] | |
stack = [(x, y)] | |
region = [] | |
while stack: | |
cx, cy = stack.pop() | |
if (cx, cy) in visited: | |
continue | |
visited.add((cx, cy)) | |
region.append((cx, cy)) | |
for nx, ny in get_neighbors(cx, cy, len(grid), len(grid[0])): | |
if grid[nx][ny] == plant_type and (nx, ny) not in visited: | |
stack.append((nx, ny)) | |
return region | |
def calculate_perimeter(region, grid): | |
perimeter = 0 | |
for x, y in region: | |
for nx, ny in get_neighbors(x, y, len(grid), len(grid[0])): | |
if grid[nx][ny] != grid[x][y]: | |
perimeter += 1 | |
return perimeter | |
def calculate_sides(region, grid): | |
sides = set() | |
for x, y in region: | |
for nx, ny in get_neighbors(x, y, len(grid), len(grid[0])): | |
if grid[nx][ny] != grid[x][y]: | |
sides.add(((x, y), (nx, ny))) | |
return len(sides) | |
def calculate_total_price(grid, use_sides=False): | |
visited = set() | |
total_price = 0 | |
for x in range(len(grid)): | |
for y in range(len(grid[0])): | |
if (x, y) not in visited: | |
region = flood_fill(grid, x, y, visited) | |
area = len(region) | |
if use_sides: | |
sides = calculate_sides(region, grid) | |
price = area * sides | |
else: | |
perimeter = calculate_perimeter(region, grid) | |
price = area * perimeter | |
total_price += price | |
return total_price | |
def main(): | |
grid = read_input(file) | |
# Part 1 | |
total_price_part1 = calculate_total_price(grid, use_sides=False) | |
print(total_price_part1) | |
# Part 2 | |
total_price_part2 = calculate_total_price(grid, use_sides=True) | |
print(total_price_part2) | |
main() |