Advent of code 2021 day 19

def process1(s, part2=False):
  scanners = []
  for i, s2 in enumerate(s.strip().split('\n\n')):
    lines = s2.split('\n')
    array = np.array([list(map(int, line.split(','))) for line in lines[1:]])
    scanners.append(array)

  scanner_transforms = [{i: np.eye(4, dtype=int)} for i in range(len(scanners))]
  scanner_rep = list(range(len(scanners)))
  TRANSFORMS = tuple(
      np.array(matrix)
      for diag in itertools.product([-1, 1], repeat=3)
      for matrix in itertools.permutations(np.diag(diag))
  )

  def compute_signatures():
    all_signatures = []
    for scanner in scanners:
      n = len(scanner)
      comb = itertools.chain.from_iterable(itertools.combinations(range(n), 2))
      indices = np.fromiter(comb, dtype=int, count=n * (n - 1))
      points = scanner[indices].reshape(-1, 2, 3)
      diff = np.sort(abs(points[:, 1] - points[:, 0]), axis=-1)
      signatures = dict(zip((tuple(e) for e in diff),
                            (tuple(e) for e in indices.reshape(-1, 2))))
      all_signatures.append(signatures)
    return all_signatures

  all_signatures = compute_signatures()
  all_signature_sets = [set(signature) for signature in all_signatures]
  intersection_counts = [
      (len(all_signature_sets[i] & all_signature_sets[j]), i, j)
      for i, j in itertools.combinations(range(len(scanners)), 2)]

  for _, i, j in sorted(intersection_counts, reverse=True):
    ir, jr = scanner_rep[i], scanner_rep[j]
    if ir == jr:
      continue  # Already joined.
    intersection = all_signature_sets[i] & all_signature_sets[j]

    match_count = collections.defaultdict(lambda: collections.defaultdict(int))
    for encoding, indices_i in all_signatures[i].items():
      if encoding in intersection:
        indices_j = all_signatures[j][encoding]
        for index_i in indices_i:
          for index_j in indices_j:
            match_count[index_i][index_j] += 1

    index_mapping = {}
    for index_i in range(len(scanners[i])):
      max_count, index_j = max(
          ((count, index) for index, count in match_count[index_i].items()),
          default=(0, 0))
      if max_count >= 11:  # Heuristically selected.
        index_mapping[index_i] = index_j

    def get_transform4():
      index_pairs = list(index_mapping.items())
      index_i0, index_j0 = index_pairs[0]
      for transform in TRANSFORMS:
        offset = scanners[i][index_i0] - (transform @ scanners[j][index_j0])
        for index_i, index_j in index_pairs[1:]:
          transformed = (transform @ scanners[j][index_j]) + offset
          if np.any(transformed != scanners[i][index_i]):
            break
        else:
          return np.vstack((np.hstack((transform, offset[:, None])),
                            [[0, 0, 0, 1]]))

    transform4 = get_transform4()  # To i from j.  (We want ir from jr.)
    transform4 = scanner_transforms[ir][i] @ transform4 @ np.linalg.inv(
        scanner_transforms[jr][j].astype(np.float32)).astype(int)

    points_ir = {tuple(point) for point in scanners[ir]}
    new_points = []
    for point in scanners[jr]:
      transformed = (transform4 @ [*point, 1])[:3]
      if tuple(transformed) not in points_ir:
        new_points.append(transformed)
    scanners[ir] = np.concatenate((scanners[ir], new_points), axis=0)

    for j2, transform2 in scanner_transforms[jr].items():
      scanner_transforms[ir][j2] = transform4 @ transform2
      scanner_rep[j2] = ir
    scanner_rep[jr] = ir
    last_merged = ir

  if not part2:
    return len(scanners[last_merged])

  final_transforms = scanner_transforms[last_merged].values()
  return max(abs(t1[:3, 3] - t2[:3, 3]).sum()
             for t1, t2 in itertools.combinations(final_transforms, 2))