-
Notifications
You must be signed in to change notification settings - Fork 0
/
day17.py
123 lines (101 loc) · 2.67 KB
/
day17.py
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
"""
Day 17: Clumsy Crucible
"""
import heapq
from enum import IntEnum
SAMPLE_INPUT_1, SAMPLE_INPUT_2 = (
"""2413432311323
3215453535623
3255245654254
3446585845452
4546657867536
1438598798454
4457876987766
3637877979653
4654967986887
4564679986453
1224686865563
2546548887735
4322674655533
""",
"""111111111111
999999999991
999999999991
999999999991
999999999991
""",
)
class _Direction(IntEnum):
U = 0
L = 1
D = 2
R = 3
def __add__(self, other):
return _Direction((self.value + other) % len(_Direction))
def __sub__(self, other):
return _Direction((self.value - other) % len(_Direction))
# pylint: disable=too-many-locals
def _solve(data, ok, turns):
data = [[int(char) for char in line] for line in data.splitlines() if line]
queue = [(0, (0, 0, _Direction.R, 0))]
best = {(0, 0, _Direction.R, 0): 0}
while queue:
cost, state = heapq.heappop(queue)
y1, x1, direction1, distance1 = state
if y1 == len(data) - 1 and x1 == len(data[-1]) - 1 and ok(distance1):
return cost
if cost > best[state]:
continue
for direction2 in turns(direction1, distance1):
y2, x2 = y1, x1
match direction2:
case _Direction.U:
y2 -= 1
case _Direction.L:
x2 -= 1
case _Direction.D:
y2 += 1
case _Direction.R:
x2 += 1
if not (0 <= y2 < len(data) and 0 <= x2 < len(data[y2])):
continue
cost2 = cost + data[y2][x2]
distance2 = distance1 + 1 if direction1 == direction2 else 1
state2 = (y2, x2, direction2, distance2)
if state2 in best and best[state2] <= cost2:
continue
best[state2] = cost2
heapq.heappush(queue, (cost2, state2))
return None
# pylint: disable=unused-argument
def _part1_ok(distance):
return True
def _part1_turns(direction, distance):
turns = [direction - 1, direction + 1]
if distance < 3:
turns.append(direction)
return turns
def part1(data):
"""
>>> part1(SAMPLE_INPUT_1)
102
"""
return _solve(data, _part1_ok, _part1_turns)
def _part2_ok(distance):
return distance >= 4
def _part2_turns(direction, distance):
turns = []
if distance >= 4:
turns.extend([direction - 1, direction + 1])
if distance < 10:
turns.append(direction)
return turns
def part2(data):
"""
>>> part2(SAMPLE_INPUT_1)
94
>>> part2(SAMPLE_INPUT_2)
71
"""
return _solve(data, _part2_ok, _part2_turns)
parts = (part1, part2)