aoc2023/day10.py at main · Liedtke/aoc2023 · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
dirs = [
    [-1, 0], # West
    [0, -1], # North
    [1, 0],  # East
    [0, 1],  # South
]
tiles = {
    #     W, N, E, S
    ".": [0, 0, 0, 0],
    "F": [0, 0, 1, 1],
    "|": [0, 1, 0, 1],
    "L": [0, 1, 1, 0],
    "7": [1, 0, 0, 1],
    "-": [1, 0, 1, 0],
    "J": [1, 1, 0, 0],
}

grid = open("input/day10.txt").read().split("\n")

def move(coord, dir):
    return (coord[0] + dir[0], coord[1] + dir[1])

def get_tile(coord):
    return grid[coord[1]][coord[0]]

def next_dir_idx(coord, current_dir_idx):
    tile = get_tile(coord)
    dir_idx = (current_dir_idx + 3) % 4
    while tiles[tile][dir_idx] == 0:
        dir_idx = (dir_idx + 1) % 4
    return dir_idx

for i, line in enumerate(grid):
    try:
        j = line.index("S")
        if j != -1:
            start = j, i
            break
    except:
        pass

# Find first tile connected to S.
assert get_tile(start) == "S"
for dir_idx, dir in enumerate(dirs):
    coord = move(start, dir)
    try:
        if tiles[get_tile(coord)][(dir_idx + 2) % 4]:
            break
    except:
        pass

# The loop mask is used for part 2 to color everything outside the loop.
loop_mask = [
    [0 for _ in range(len(grid[0]) * 2 + 1)]
    for _ in range(len(grid) * 2 + 1)]

def mark(coord, dir_idx):
    x = coord[0] * 2 + 1
    y = coord[1] * 2 + 1
    loop_mask[y][x] = 1
    dir = dirs[dir_idx]
    x = x - dir[0]
    y = y - dir[1]
    loop_mask[y][x] = 1

# Follow pipes until S.
mark(coord, dir_idx)
steps = 1
while coord != start:
    dir_idx = next_dir_idx(coord, dir_idx)
    coord = move(coord, dirs[dir_idx])
    mark(coord, dir_idx)
    steps += 1

# Color everything outside the loop.
coord = (0, 0)
to_color = [coord]
visited = set()
while to_color:
    coord = to_color.pop()
    try:
        if loop_mask[coord[1]][coord[0]] == 0:
            loop_mask[coord[1]][coord[0]] = 2
            for dir in dirs:
                next = move(coord, dir)
                if next not in visited:
                    visited.add(next)
                    to_color.append(next)
    except:
        pass

# Iterate over the grid and check if the coord is marked as inside the loop.
res_b = sum(sum(loop_mask[y * 2 + 1][x * 2 + 1] == 0 for x in range(len(grid[0])))
          for y in range(len(grid)))
print(f"a = {steps // 2}")
print(f"b = {res_b}")