day19_part1.py (10677B)
1 #!/usr/bin/env python 2 """Advent of Code 2021, day 19 (part 1): Beacon Scanner 3 Given positions (but not consistent frame of reference), find the overlap 4 between a set of beacons detected by several scanners""" 5 6 # Looking at the leaderboard, this one has been hard for folks. My first 7 # impression is that its pretty doable if I can get away with brute forcing a 8 # few parts of it, but if I need to be sophisticated (i.e., if a bunch of 9 # polynomial time approaches dont scale), then this could be a huge slog. 10 # 11 # Numpy and SciPy are helping out here, but neither are doing any algorithmic 12 # heavy lifting (as was the case when using optimization or graph search from 13 # SciPy in previous puzzles). 14 # 15 # Note: I'm using _augmented matrices_ to represent the beacon locations and 16 # transformations. 3D data is represented by 4D vectors padded with a 1, 17 # allowing translations to be implemented with matrix multiplication. 18 19 from typing import List, Union, Tuple 20 from itertools import product 21 import numpy as np 22 from scipy.spatial.transform import Rotation as R 23 from scipy.spatial.distance import cdist 24 from utils import get_puzzle_input 25 26 EXAMPLE_INPUT = \ 27 """--- scanner 0 --- 28 404,-588,-901 29 528,-643,409 30 -838,591,734 31 390,-675,-793 32 -537,-823,-458 33 -485,-357,347 34 -345,-311,381 35 -661,-816,-575 36 -876,649,763 37 -618,-824,-621 38 553,345,-567 39 474,580,667 40 -447,-329,318 41 -584,868,-557 42 544,-627,-890 43 564,392,-477 44 455,729,728 45 -892,524,684 46 -689,845,-530 47 423,-701,434 48 7,-33,-71 49 630,319,-379 50 443,580,662 51 -789,900,-551 52 459,-707,401 53 54 --- scanner 1 --- 55 686,422,578 56 605,423,415 57 515,917,-361 58 -336,658,858 59 95,138,22 60 -476,619,847 61 -340,-569,-846 62 567,-361,727 63 -460,603,-452 64 669,-402,600 65 729,430,532 66 -500,-761,534 67 -322,571,750 68 -466,-666,-811 69 -429,-592,574 70 -355,545,-477 71 703,-491,-529 72 -328,-685,520 73 413,935,-424 74 -391,539,-444 75 586,-435,557 76 -364,-763,-893 77 807,-499,-711 78 755,-354,-619 79 553,889,-390 80 81 --- scanner 2 --- 82 649,640,665 83 682,-795,504 84 -784,533,-524 85 -644,584,-595 86 -588,-843,648 87 -30,6,44 88 -674,560,763 89 500,723,-460 90 609,671,-379 91 -555,-800,653 92 -675,-892,-343 93 697,-426,-610 94 578,704,681 95 493,664,-388 96 -671,-858,530 97 -667,343,800 98 571,-461,-707 99 -138,-166,112 100 -889,563,-600 101 646,-828,498 102 640,759,510 103 -630,509,768 104 -681,-892,-333 105 673,-379,-804 106 -742,-814,-386 107 577,-820,562 108 109 --- scanner 3 --- 110 -589,542,597 111 605,-692,669 112 -500,565,-823 113 -660,373,557 114 -458,-679,-417 115 -488,449,543 116 -626,468,-788 117 338,-750,-386 118 528,-832,-391 119 562,-778,733 120 -938,-730,414 121 543,643,-506 122 -524,371,-870 123 407,773,750 124 -104,29,83 125 378,-903,-323 126 -778,-728,485 127 426,699,580 128 -438,-605,-362 129 -469,-447,-387 130 509,732,623 131 647,635,-688 132 -868,-804,481 133 614,-800,639 134 595,780,-596 135 136 --- scanner 4 --- 137 727,592,562 138 -293,-554,779 139 441,611,-461 140 -714,465,-776 141 -743,427,-804 142 -660,-479,-426 143 832,-632,460 144 927,-485,-438 145 408,393,-506 146 466,436,-512 147 110,16,151 148 -258,-428,682 149 -393,719,612 150 -211,-452,876 151 808,-476,-593 152 -575,615,604 153 -485,667,467 154 -680,325,-822 155 -627,-443,-432 156 872,-547,-609 157 833,512,582 158 807,604,487 159 839,-516,451 160 891,-625,532 161 -652,-548,-490 162 30,-46,-14 163 """ 164 165 def parse_input(input_string: str) -> List[np.ndarray]: 166 """Given the puzzle input, return a list comprising all the 167 beacons detected by each scanner""" 168 # I do not love how hacky this is 169 scanner_beacons = [] 170 beacon_start = 0 171 # Depends on the trailing newline in the puzzle input. Hacky! 172 input_lines = input_string.split('\n') 173 for idx, line in enumerate(input_lines): 174 if line.startswith('---'): 175 beacon_start = idx + 1 176 elif line == '': 177 # End of a probe's beacons 178 scanner_beacons.append( 179 np.array([l.split(',') + ['1'] for l in input_lines[beacon_start:idx]]) 180 .astype(float) 181 ) 182 return scanner_beacons 183 184 def augment_matrix(rotation_matrix: np.ndarray) -> np.ndarray: 185 """Given a 3x3 transformation matrix, return a 4x4 augmented matrix 186 representing the same transformation""" 187 return np.vstack(( 188 np.hstack((rotation_matrix, np.zeros((3, 1)))), 189 np.array([0, 0, 0, 1]) 190 )) 191 192 def create_rotations() -> List[np.ndarray]: 193 """Create the 24 distinct rotation matrices for a probe that can be 194 oriented to point along any direction of the x, y, and z axis, with its top 195 aligned along one of the other two axes""" 196 # This function is inefficient; it considers all possible 90 degree 197 # rotations along each of the three axes (i.e., 4^3 = 64 combinations), and 198 # then only returns the 24 unique ones. One could probably reason out how 199 # to more directly generate the 24 that matter, but this works! 200 # This first step creates a Rotation instance that stacks all 64 rotations 201 rotations = R.from_euler( 202 'zyx', 203 np.array(list(product(*(((0, 90, 180, 270),) * 3)))), 204 degrees=True 205 ) 206 # This next step creates a list of non-duplicate rotation matricies 207 final_rotations = [augment_matrix( 208 rotations[0].as_matrix().round() 209 )] 210 for rot_obj in rotations[1:]: 211 # In our case, all of the elements of the rotation matrix should be -1, 212 # 0, or 1, so we'll round everything to the nearest value to make 213 # equality checking easier 214 rot_mat = augment_matrix(rot_obj.as_matrix().round()) 215 for prev_rot_mat in final_rotations: 216 if (rot_mat == prev_rot_mat).all(): 217 break 218 else: 219 # We get here if the for loop completed without breaking, which 220 # means we have a unique rotation matrix 221 final_rotations.append(rot_mat) 222 return final_rotations 223 224 def translation_matrix(offset: np.ndarray) -> np.ndarray: 225 """Given a length-three array representing the x, y, and z offset, return a 226 4x4 translation matrix (if the offset array is longer than three elements, 227 only the first three are used, so it's safe to lazily pass length-four 228 arrays too)""" 229 return np.vstack(( 230 np.identity(4)[0:3, :], 231 np.append(offset[0:3], 1) 232 )) 233 234 def count_overlapping_beacons(scanner_1: np.ndarray, 235 scanner_2: np.ndarray) -> int: 236 """Given a pair of (appropriately rotated and translated) beacon positions, 237 count how many from each scanner are at approximately the same location""" 238 return (cdist(scanner_1, scanner_2) < 1e-8).sum() 239 240 def find_best_overlap(scanner_1: np.ndarray, 241 scanner_2: np.ndarray) -> Union[np.ndarray, None]: 242 """Find the combined rotation/translation matrix for scanner 2 that results 243 in (magic number) 12 overlapping beacons, or return None if no such 244 transformation exists""" 245 # The algorithm: for each possible rotation, calculate the distance matrix 246 # (using Euclidean distance) between the two sets of beacons. If (magic 247 # number) 12 or more beacon pairs have the same distance, translate 248 # scanner_2's beacons so that the first such pair has a distance of zero. 249 # If the remainder of the same-distance pairs don't also go to zero, 250 # continue to the next pair, until there are fewer than twelve remaining 251 # pairs to check. 252 for rot_mat in create_rotations(): 253 scanner_2_rot = scanner_2.dot(rot_mat) 254 dist_mat = cdist(scanner_1, scanner_2_rot) 255 # We'll round the distance matrix to get a sense of the mode 256 beacon_dist_count = np.bincount( 257 np.reshape(dist_mat.round(), dist_mat.size) 258 .astype(int) 259 ) 260 if beacon_dist_count.max() >= 12: 261 # There are at least 12 beacon pairs with similar distances. 262 # Iterate through the pairs to see if translating one to the other 263 # provides a match. 264 potential_beacon_matches = np.nonzero( 265 dist_mat.round() == np.argmax(beacon_dist_count) 266 ) 267 # Technically we can stop iterating when there are fewer than 12 268 # possible matches but I don't feel like programming that in unless 269 # I have to. 270 for scanner_1_beacon_idx, scanner_2_beacon_idx in \ 271 zip(*potential_beacon_matches): 272 scanner_1_beacon = scanner_1[scanner_1_beacon_idx, :] 273 scanner_2_beacon = scanner_2_rot[scanner_2_beacon_idx, :] 274 trans_mat = translation_matrix( 275 scanner_1_beacon - scanner_2_beacon 276 ) 277 num_overlapping = count_overlapping_beacons( 278 scanner_1, 279 scanner_2_rot.dot(trans_mat) 280 ) 281 if num_overlapping >= 12: 282 return rot_mat.dot(trans_mat) 283 return None 284 285 def create_beacon_network(scanner_beacons: List[np.ndarray]) -> \ 286 Tuple[np.ndarray, List[np.ndarray]]: 287 """Given a list of scanner beacons, find all of the unique beacons across 288 all scanners by rotating and translating the scanner coordinates (relative 289 to the first beacon in the list). Returns the final beacon locations and a 290 list of translation matrices.""" 291 scanner_beacon_list: Union[None, np.ndarray] = scanner_beacons.copy() 292 translation_list: Union[None, np.ndarray] = [None] * len(scanner_beacon_list) 293 294 final_beacons = scanner_beacon_list[0] 295 scanner_beacon_list[0] = None 296 translation_list[0] = np.identity(4) 297 298 # The way the problem is written suggests that not every pair of scanners 299 # has an overlapping set of beacons. This algorithm loops through all 300 # scanners sequentially. When it finds overlap with the network, its 301 # beacons are added to the network and it's removed from further 302 # consideration. Eventually, all of the beacons should get added! 303 while sum([x is not None for x in scanner_beacon_list]) > 0: 304 for scanner_id, beacons in enumerate(scanner_beacon_list): 305 if beacons is not None: 306 trans_mat = find_best_overlap(final_beacons, beacons) 307 if trans_mat is not None: 308 final_beacons = np.unique( 309 np.vstack(( 310 final_beacons, 311 beacons.dot(trans_mat) 312 )), 313 axis=0 314 ) 315 316 scanner_beacon_list[scanner_id] = None 317 translation_list[scanner_id] = trans_mat 318 319 return final_beacons, translation_list 320 321 def solve_puzzle(input_string: str) -> int: 322 """Return the numeric solution to the puzzle""" 323 return create_beacon_network(parse_input(input_string))[0].shape[0] 324 325 def main() -> None: 326 """Run when the file is called as a script""" 327 assert solve_puzzle(EXAMPLE_INPUT) == 79 328 print("Total number of beacons:", 329 solve_puzzle(get_puzzle_input(19))) 330 331 if __name__ == "__main__": 332 main()