aoc2021-19.py

import itertools
import functools
import sys

AXIS_ORIENTATIONS = list(itertools.product([-1, 1], [-1, 1], [-1, 1]))
AXIS_ORDERS = list(itertools.permutations(range(3), 3))


def shift_scanner(scanner, vector):
    return tuple(tuple(c + d for c, d in zip(point, vector)) for point in scanner)


@functools.lru_cache(maxsize=None)
def get_mutants(scanner):
    def iter_mutants(scanner):
        for orientation in AXIS_ORIENTATIONS:
            # apply axis negation
            mutant = tuple(
                tuple(c * o for c, o in zip(point, orientation)) for point in scanner
            )
            for order in AXIS_ORDERS:
                yield tuple(
                    tuple(point[order[i]] for i in range(3)) for point in mutant
                )

    return list(iter_mutants(scanner))


@functools.lru_cache(maxsize=None)
def align(target, original_scanner):
    for scanner in get_mutants(original_scanner):
        for target_point, scanner_point in itertools.product(target, scanner):
            delta = tuple(target_point[i] - scanner_point[i] for i in range(3))
            shifted_scanner = shift_scanner(scanner, delta)
            matches = len(set(shifted_scanner) & set(target))
            if matches >= 12:
                # match!
                return shifted_scanner, delta
    # no match :(
    return None, None


SCANNERS = [
    tuple(
        tuple(int(x.strip()) for x in l.split(","))
        for l in s.split("\n")
        if len(l.split(",")) == 3
    )
    for s in sys.stdin.read().split("\n\n")
]

unaligned = list(SCANNERS)
initial = unaligned.pop()
aligned = [initial]
positions = {initial: (0, 0, 0)}


while unaligned:
    for target, scanner in itertools.product(aligned, unaligned):
        match, delta = align(target, scanner)
        if match:
            unaligned.remove(scanner)
            aligned.append(match)
            positions[match] = delta
            break
    else:
        assert False, "No solution"

print("Part 1:", len(set(sum([list(scanner) for scanner in aligned], start=[]))))

print(
    "Part 2:",
    max(
        sum(abs(c1 - c2) for c1, c2 in zip(sp1, sp2))
        for sp1, sp2 in itertools.product(positions.values(), positions.values())
    ),
)