AOC2021 Day 19 (Python)

day19_tests.py

def read_file(file):
    stripped = (line.strip() for line in file)
    scanners = []

    for line in stripped:
        if line.startswith('---'):
            scanner = set()
            scanners.append(scanner)
        elif len(line) > 0:
            x_str, y_str, z_str = line.split(',')
            scanner.add((int(x_str), int(y_str), int(z_str)))

    return scanners


def transform(view, transformation):
    return set(transformation(x, y, z) for (x, y, z) in view)


identity_matrix = (
    (1, 0, 0),
    (0, 1, 0),
    (0, 0, 1)
)

rotate_90_xy = (
    (0, -1, 0),
    (1, 0, 0),
    (0, 0, 1)
)

rotate_90_yz = (
    (1, 0, 0),
    (0, 0, -1),
    (0, 1, 0)
)

rotate_90_xz = (
    (0, 0, -1),
    (0, 1, 0),
    (1, 0, 0)
)


def matrix_mul(p, q):
    return tuple(
        tuple(sum(e * q[i][col_index] for i, e in enumerate(row))
              for col_index in range(len(q[0])))
        for row in p)


rotate_180_xy = matrix_mul(rotate_90_xy, rotate_90_xy)
rotate_270_xy = matrix_mul(rotate_180_xy, rotate_90_xy)

xy_rotations = [identity_matrix, rotate_90_xy, rotate_180_xy, rotate_270_xy]

rotate_180_yz = matrix_mul(rotate_90_yz, rotate_90_yz)
rotate_270_yz = matrix_mul(rotate_180_yz, rotate_90_yz)

yz_rotations = [identity_matrix, rotate_90_yz, rotate_180_yz, rotate_270_yz]

rotate_180_xz = matrix_mul(rotate_90_xz, rotate_90_xz)
rotate_270_xz = matrix_mul(rotate_180_xz, rotate_90_xz)

xz_rotations = [identity_matrix, rotate_90_xz, rotate_180_xz, rotate_270_xz]

all_rotations = set(matrix_mul(c, matrix_mul(a, b)) for a in xy_rotations for b in yz_rotations for c in xz_rotations)

assert len(all_rotations) == 24


def view_rotations(view):
    return (transform(view, lambda x, y, z: matrix_mul(((x, y, z), ), rotation)[0])
            for rotation in all_rotations)


def get_offset(a_beacon, b_beacon):
    ax, ay, az = a_beacon
    bx, by, bz = b_beacon
    return bx - ax, by - ay, bz - az


def find_intersection(a_beacons, b_beacons):
    candidate_offsets = dict()
    for a_beacon in a_beacons:
        for b_beacon in b_beacons:
            offset = get_offset(a_beacon, b_beacon)
            pairs = candidate_offsets.get(offset, [])
            pairs.append((a_beacon, b_beacon))
            candidate_offsets[offset] = pairs

    intersection_offset, matches = max(candidate_offsets.items(), key=lambda i: len(i[1]))
    return intersection_offset, len(matches)


def apply_offset(a, b):
    ax, ay, az = a
    bx, by, bz = b
    return ax + bx, ay + by, az + bz


def scanners_agree(scanner_a_view, scanner_b_view):
    for rotation in view_rotations(scanner_b_view):
        offset, count = find_intersection(scanner_a_view, rotation)
        if count >= 12:
            offset_x, offset_y, offset_z = offset
            scanner_b_absolute_position = (-offset_x, -offset_y, -offset_z)

            absolute_beacon_positions = set(apply_offset(position, scanner_b_absolute_position)
                                            for position in rotation)
            return True, scanner_b_absolute_position, absolute_beacon_positions

    return False, None, None


def agreeing_pairs(scanners):
    start, *remainder = scanners
    scanner_positions = [(0, 0, 0)]
    fixed_beacons = start

    unmatched = set(range(len(remainder)))
    while len(unmatched) > 0:
        to_remove = set()

        for idx in unmatched:
            scanner = remainder[idx]
            agree, position, beacon_positions = scanners_agree(fixed_beacons, scanner)
            if agree:
                scanner_positions.append(position)
                fixed_beacons = fixed_beacons.union(beacon_positions)
                to_remove.add(idx)

        unmatched = unmatched.difference(to_remove)

    return scanner_positions, fixed_beacons


def test_agreeing_pairs():
    with open("puzzle_inputs/day19.txt") as file:
        scanners = read_file(file)

    scanner_positions, fixed_beacons = agreeing_pairs(scanners)

    print()
    print(len(fixed_beacons))

    max_scanner_distance = max(abs(ax - bx) + abs(ay - by + abs(az - bz))
                               for (ax, ay, az) in scanner_positions
                               for (bx, by, bz) in scanner_positions)

    print(max_scanner_distance)