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
|
from collections import Counter, defaultdict
import numpy as np
def fit_homography(P1, P2):
# p is a size (N, 3) set of 3d points
# X is a size (N, 4) set of 3d points
N, _ = P1.shape
M = np.zeros((N*3, 12))
b = np.zeros((N*3, 1))
for i, (p, y) in enumerate(zip(P1, P2)):
b[i*3:i*3+3, 0] = y
M[i*3, :3] = p
M[i*3 + 1, 3:6] = p
M[i*3 + 2, 6:9] = p
M[i*3: i*3 + 3, 9:] = np.eye(3)
x = np.linalg.inv(M.T @ M) @ M.T @ b
R = x[:9].reshape(3,3).round().astype(int)
t = x[9:].round().astype(int)
return R, t
class SensorData:
def __init__(self, points: np.ndarray):
self.points = points
self.point_to_dists = defaultdict(set)
self.sensors = np.zeros((1,3))
# Get dists for points
self._find_dists()
def _find_dists(self):
self.point_to_dists = defaultdict(set)
for p1 in self.points:
for p2 in self.points:
if tuple(p1) == tuple(p2): continue
d = ((p1 - p2) ** 2).sum().item()
self.point_to_dists[tuple(p1)].add(d)
self.point_to_dists[tuple(p2)].add(d)
def add_points(self, other):
matches = self.match_points(other)
if matches is None or matches.shape[0] < 12:
return False
P1, P2 = self.points[matches[:,0], :], other.points[matches[:,1], :]
R, t = fit_homography(P2, P1)
transformed = other.points @ R.T + t.T
self.points = np.concatenate([self.points, transformed])
self.points = np.unique(self.points, axis=0)
# self._find_dists()
for new_point, old_point in zip(transformed, other.points):
self.point_to_dists[tuple(new_point)].update(other.point_to_dists[tuple(old_point)])
# Also map the sensor location to store list of sensors
self.sensors = np.concatenate([
self.sensors,
other.sensors @ R.T + t.T
])
return True
def match_points(self, other):
paired = []
for j, p2 in enumerate(other.points):
for i, p1 in enumerate(self.points):
dists1 = self.point_to_dists[tuple(p1)]
dists2 = other.point_to_dists[tuple(p2)]
common = dists1.intersection(dists2)
if len(common) < 11: continue
paired.append((i, j))
if len(paired) == 12:
break
if not paired: return None
return np.stack(paired)
def max_dist(self):
best = 0
for s1 in self.sensors:
for s2 in self.sensors:
d = np.abs(s1 - s2).sum().item()
best = max(d, best)
return int(best)
def parse_input(contents):
sensors = contents.split('\n\n')
output = []
for sensor in sensors:
output.append(
np.stack(
[list(map(int, line.split(',')))
for line in sensor.splitlines()[1:]]
)
)
return output
def part_1(sensors):
sensors = [SensorData(np.copy(n)) for n in sensors]
data = sensors[0]
to_pair = set(range(1, len(sensors)))
while to_pair:
again = set()
for i in to_pair:
if data.add_points(sensors[i]): continue
again.add(i)
to_pair = again
return data
if __name__ == "__main__":
print('--- Part 1 ---')
with open('day19.txt') as f:
sensors = parse_input(f.read())
data = part_1(sensors)
print(data.points.shape[0])
print('--- Part 2 ---')
print(data.max_dist())
|