aoc

Advent of Code
git clone https://git.pngu.org/aoc
Log | Files | Refs | README

16_2024.py (1898B)

 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
import heapq

input = "data/16_2024.txt"
data = open(input).read().strip()
lines = data.split("\n")
rows = len(lines)
cols = len(lines[0])
lines = [[lines[r][c] for c in range(cols)] for r in range(rows)]
dirs = [(-1, 0), (0, 1), (1, 0), (0, -1)]


for r in range(rows):
    for c in range(cols):
        if lines[r][c] == "S":
            sr, sc = r, c
        if lines[r][c] == "E":
            er, ec = r, c


q = []
vis = set()
heapq.heappush(q, (0, sr, sc, 1))
dist = {}
best = None
while q:
    d, r, c, dir = heapq.heappop(q)
    if (r, c, dir) not in dist:
        dist[(r, c, dir)] = d
    if r == er and c == ec and best is None:
        best = d
    if (r, c, dir) in vis:
        continue
    vis.add((r, c, dir))
    dr, dc = dirs[dir]
    rr, cc = r + dr, c + dc
    if 0 <= cc < cols and 0 <= rr < rows and lines[rr][cc] != "#":
        heapq.heappush(q, (d + 1, rr, cc, dir))
    heapq.heappush(q, (d + 1000, r, c, (dir + 1) % 4))
    heapq.heappush(q, (d + 1000, r, c, (dir + 3) % 4))


print(best)


q = []
vis = set()
dist2 = {}
for dir in range(4):
    heapq.heappush(q, (0, er, ec, dir))


while q:
    d, r, c, dir = heapq.heappop(q)
    if (r, c, dir) not in dist2:
        dist2[(r, c, dir)] = d
    if (r, c, dir) in vis:
        continue
    vis.add((r, c, dir))
    dr, dc = dirs[(dir + 2) % 4]
    rr, cc = r + dr, c + dc
    if 0 <= cc < cols and 0 <= rr < rows and lines[rr][cc] != "#":
        heapq.heappush(q, (d + 1, rr, cc, dir))
    heapq.heappush(q, (d + 1000, r, c, (dir + 1) % 4))
    heapq.heappush(q, (d + 1000, r, c, (dir + 3) % 4))


valid = set()
for r in range(rows):
    for c in range(cols):
        for dir in range(4):
            if (
                (r, c, dir) in dist
                and (r, c, dir) in dist2
                and dist[(r, c, dir)] + dist2[(r, c, dir)] == best
            ):
                valid.add((r, c))


print(len(valid))