-
Notifications
You must be signed in to change notification settings - Fork 1
/
solve.coco
107 lines (74 loc) · 2.27 KB
/
solve.coco
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
from pathlib import Path
from typing import TextIO
import numpy as np
def get_inp(fp: TextIO):
d = {"O": 2, "#": 1, ".": 0}
return np.array([[d[x] for x in line.strip()] for line in fp])
def pre(x):
i = 0
for c in x:
if c == 0:
i += 1
else:
break
return i
def north(mat):
for i, j in np.argwhere(mat == 2):
up = mat[:i, j][::-1]
if up.size == 0 or up[0] != 0:
continue
fst = pre(up) # TODO: there must be a better way
# print(i, j, '->>', up, fst, '->>', i-fst, j)
mat[i, j] = 0 # empty the current cell
mat[i - fst, j] = 2 # north
return mat
def south(mat) =
mat |> np.rot90$(k=2) |> north |> np.rot90$(k=-2)
def west(mat) =
mat |> np.rot90$(k=-1) |> north |> np.rot90$(k=1)
def east(mat) =
mat |> np.rot90$(k=1) |> north |> np.rot90$(k=-1)
def single_cycle(mat) =
mat |> north |> west |> south |> east
def load(mat):
return (len(mat) - np.argwhere(mat == 2)[:, 0]).sum()
def iterate(f, x):
y = f(x)
yield y
yield from iterate(f, y)
def multiple_cycles(mat, n):
seen = {}
for i, mat in enumerate(iterate(single_cycle, mat), start=1):
h = mat |> map$(tuple ..> hash) |> tuple |> hash
if h not in seen:
# save the arrays until a period is found to avoid simulating again
seen[h] = (i, mat.copy())
else:
# (initial) + | ... | + | ... | + ...
# ^ --- ^ period size
j, _ = seen[h]
initial = j - 1
period_size = i - j
d, m = divmod(n - initial, period_size)
assert n == initial + d * period_size + m
for i, mat in seen.values():
if i == initial + m:
return mat
assert False
def solve(fp: TextIO):
mat = get_inp(fp)
p1 = mat.copy() |> north |> load
p2 = mat.copy() |> multiple_cycles$(n=int(1e9)) |> load
return p1, p2
def test_example():
with open(Path(__file__).parent / "example") as fp:
p1, p2 = solve(fp)
assert p1 == 136
assert p2 == 64
def main():
with open(0) as fp:
p1, p2 = solve(fp)
print(f"Part 1: {p1}")
print(f"Part 2: {p2}")
if __name__ == "__main__":
main()