advent-of-code-2021/src/day19.py at master · markusremplbauer/advent-of-code-2021 · GitHub

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
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
# Advent of Code 2021, Day 19
# markusremplbauer

from collections import defaultdict, deque
from functools import lru_cache
from itertools import product

from aocd.models import Puzzle
from funcy import collecting, print_calls
from parse import parse
from tqdm import tqdm

ROTATIONS = [
    lambda x, y, z: (-x, -y, z),
    lambda x, y, z: (-x, -z, -y),
    lambda x, y, z: (-x, y, -z),
    lambda x, y, z: (-x, z, y),
    lambda x, y, z: (-y, -x, -z),
    lambda x, y, z: (-y, -z, x),
    lambda x, y, z: (-y, x, z),
    lambda x, y, z: (-y, z, -x),
    lambda x, y, z: (-z, -x, y),
    lambda x, y, z: (-z, -y, -x),
    lambda x, y, z: (-z, x, -y),
    lambda x, y, z: (-z, y, x),
    lambda x, y, z: (x, -y, -z),
    lambda x, y, z: (x, -z, y),
    lambda x, y, z: (x, y, z),
    lambda x, y, z: (x, z, -y),
    lambda x, y, z: (y, -x, z),
    lambda x, y, z: (y, -z, -x),
    lambda x, y, z: (y, x, -z),
    lambda x, y, z: (y, z, x),
    lambda x, y, z: (z, -x, -y),
    lambda x, y, z: (z, -y, x),
    lambda x, y, z: (z, x, y),
    lambda x, y, z: (z, y, -x),
]


@print_calls
def part1(scans):
    reference, _ = solve(scans)
    return len(reference)


@print_calls
def part2(scans):
    _, offsets = solve(scans)
    largest = 0

    # compute Manhattan distance and store largest
    for (ax, ay, az), (bx, by, bz) in product(offsets, repeat=2):
        dist = abs(ax - bx) + abs(ay - by) + abs(az - bz)
        largest = max(largest, dist)

    return largest


@lru_cache
def solve(scans, min_beacons=12):
    reference = set(scans[0])
    queue = deque(scans.keys())
    queue.remove(0)

    offsets = []

    with tqdm(total=len(queue)) as progress:
        # match scanner by scanner
        while queue:
            i = queue.pop()

            match = match_beacons(reference, scans[i], min_beacons)
            if not match:
                # no match found, re-visit this scanner at the end again
                queue.appendleft(i)
                continue

            # transform scanner beacons and merge them with the reference
            r, offset = match
            beacons = set(transform(scans[i], r, offset))
            reference.update(beacons)

            offsets.append(offset)
            progress.update(1)

    return reference, offsets


def transform(scan, rot, offset):
    # rotate and translate all points in scan
    dx, dy, dz = offset
    for s in scan:
        sx, sy, sz = ROTATIONS[rot](*s)
        yield sx + dx, sy + dy, sz + dz


def match_beacons(scan_a, scan_b, min_beacons):
    # rotate b in all different directions and try extracting the offset
    for r, rot in enumerate(ROTATIONS):
        b_rot = rotated(scan_b, rot)
        offset = scanner_offset(scan_a, b_rot, min_beacons)
        if offset:
            return r, offset


def scanner_offset(scan_a, scan_b, min_beacons):
    dists = defaultdict(set)
    for p1 in scan_a:
        for p2 in scan_b:
            # group by distance vectors
            # for all points p1 from a and p2 from b
            d = distance(p2, p1)
            dists[d].add((p1, p2))

            # enough equal vectors identify the scanner offset
            if len(dists[d]) >= min_beacons:
                return d


def distance(a, b):
    (ax, ay, az), (bx, by, bz) = a, b
    return bx - ax, by - ay, bz - az


@collecting
def rotated(scan, rot):
    for xyz in scan:
        yield rot(*xyz)


def load(data):
    blocks = data.split("\n\n")
    scanners = {}
    for i, block in enumerate(blocks):
        lines = block.split("\n")[1:]
        xyz = frozenset(tuple(parse("{:d},{:d},{:d}", e)) for e in lines)
        scanners[i] = xyz
    return frozendict(scanners)


class frozendict(dict):
    def __hash__(self):
        # primitive dict hash, needed for lru_cache on solve()
        return hash((frozenset(self.keys()), frozenset(self.values())))


def main():
    puzzle = Puzzle(year=2021, day=19)

    ans1 = part1(load(puzzle.input_data))
    # puzzle.answer_a = ans1
    ans2 = part2(load(puzzle.input_data))
    # puzzle.answer_b = ans2


if __name__ == "__main__":
    main()