/
06.py
85 lines (71 loc) · 1.73 KB
/
06.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
import sys
import pytest
from pathlib import Path
import math
from textwrap import dedent
def solve(text: str) -> int:
lines = text.splitlines()
times = list(map(int, lines[0].split(":", maxsplit=1)[1].split()))
distances = list(map(int, lines[1].split(":", maxsplit=1)[1].split()))
total = 1
for i in range(len(times)):
time = times[i]
distance = distances[i]
i = math.ceil((time - math.sqrt((time**2) - (4 * distance))) / 2)
if (i * (time - i)) <= distance:
i += 1
max_peak = (time - 1) // 2
if time % 2 == 0:
max_peak += 1
if i < max_peak:
if time % 2 == 0:
subtotal = (max_peak - i) * 2 + 1
else:
subtotal = (max_peak - i + 1) * 2
total *= subtotal
return total
if __name__ == "__main__":
p = Path(sys.argv[1])
print(solve(p.read_text()))
def test():
text = dedent(
"""
Time: 7
Distance: 9
"""
).strip()
assert solve(text) == 4
text = dedent(
"""
Time: 8
Distance: 9
"""
).strip()
assert solve(text) == 5
text = dedent(
"""
Time: 15
Distance: 40
"""
).strip()
assert solve(text) == 8
text = dedent(
"""
Time: 30
Distance: 200
"""
).strip()
assert solve(text) == 9
text = dedent(
"""
Time: 7 15 30
Distance: 9 40 200
"""
).strip()
assert solve(text) == 288
def test_input():
p = Path("06.txt")
if not p.exists():
pytest.skip(f"{p} does not exist")
text = p.read_text()
assert solve(text) == 3317888