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
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
|
import re
import heapq
input = "data/21_2024.txt"
data = open(input).read().strip()
pad1 = ["789", "456", "123", " 0A"]
pad2 = [" ^A", "<v>"]
dp = {}
def ints(s):
return [int(x) for x in re.findall("-?\d+", s)]
def get_pad1(p1):
r, c = p1
if not (0 <= r < len(pad1) and 0 <= c < len(pad1[r])):
return None
if pad1[r][c] == " ":
return None
return pad1[r][c]
def get_pad2(p):
r, c = p
if not (0 <= r < len(pad2) and 0 <= c < len(pad2[r])):
return None
if pad2[r][c] == " ":
return None
return pad2[r][c]
def apply_pad1(p, move):
if move == "A":
return (p, get_pad1(p))
elif move == "<":
return ((p[0], p[1] - 1), None)
elif move == "^":
return ((p[0] - 1, p[1]), None)
elif move == ">":
return ((p[0], p[1] + 1), None)
elif move == "v":
return ((p[0] + 1, p[1]), None)
def apply_pad2(p, move):
if move == "A":
return (p, get_pad2(p))
elif move == "<":
return ((p[0], p[1] - 1), None)
elif move == "^":
return ((p[0] - 1, p[1]), None)
elif move == ">":
return ((p[0], p[1] + 1), None)
elif move == "v":
return ((p[0] + 1, p[1]), None)
def solve(code, pads):
start = [0, (3, 2), "A", "", ""]
q, vis = [], {}
heapq.heappush(q, start)
while q:
d, p1, p2, out, path = heapq.heappop(q)
assert p2 in ["<", ">", "v", "^", "A"]
if out == code:
return d
if not code.startswith(out):
continue
if get_pad1(p1) is None:
continue
key = (p1, p2, out)
if key in vis:
assert d >= vis[key], f"{key=} {d=} {vis[key]=}"
continue
vis[key] = d
for move in ["^", "<", "v", ">", "A"]:
new_p1 = p1
new_out = out
new_p1, output = apply_pad1(p1, move)
if output is not None:
new_out = out + output
cost_move = cost(move, p2, pads)
new_path = path
assert cost_move >= 0
heapq.heappush(q, [d + cost_move, new_p1, move, new_out, new_path])
def cost(ch, prev_move, pads):
key = (ch, prev_move, pads)
if key in dp:
return dp[key]
if pads == 0:
return 1
else:
assert ch in ["^", ">", "v", "<", "A"]
assert prev_move in ["^", ">", "v", "<", "A"]
assert pads >= 1
q = []
start_pos = {"^": (0, 1), "<": (1, 0), "v": (1, 1), ">": (1, 2), "A": (0, 2)}[
prev_move
]
heapq.heappush(q, [0, start_pos, "A", "", ""])
vis = {}
while q:
d, p, prev, out, path = heapq.heappop(q)
if get_pad2(p) is None:
continue
if out == ch:
dp[key] = d
return d
elif len(out) > 0:
continue
vis_key = (p, prev)
if vis_key in vis:
assert d >= vis[vis_key]
continue
vis[vis_key] = d
for move in ["^", "<", "v", ">", "A"]:
new_p, output = apply_pad2(p, move)
cost_move = cost(move, prev, pads - 1)
new_d = d + cost_move
new_path = path
new_out = out
if output is not None:
new_out = new_out + output
heapq.heappush(q, [new_d, new_p, move, new_out, new_path])
assert False, f"{ch=} {pads=}"
p1 = 0
p2 = 0
for line in data.split("\n"):
s1 = solve(line, 2)
s2 = solve(line, 25)
line_int = ints(line)[0]
p1 += line_int * s1
p2 += line_int * s2
print(p1)
print(p2)
|