advent-of-code/2024/day08/day8.py at master · buckley-w-david/advent-of-code · 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
from collections import defaultdict
from itertools import combinations

from aoc_utils import *  # type: ignore
from aocd import get_data

data = get_data(year=2024, day=8, block=True)


def part_one(data):
    locations = defaultdict(set)
    grid = Grid.parse(data)

    for yx, c in grid.row_major_with_index():
        if c != ".":
            locations[c].add(yx)

    antinodes = set()
    for _, positions in locations.items():
        for (ya, xa), (yb, xb) in combinations(positions, 2):
            dy = yb - ya
            dx = xb - xa

            p1 = (ya - dy, xa - dx)
            if p1 in grid:
                antinodes.add(p1)

            p2 = (yb + dy, xb + dx)
            if p2 in grid:
                antinodes.add(p2)

    return len(antinodes)


def part_two(data):
    locations = defaultdict(set)
    grid = Grid.parse(data)

    for yx, c in grid.row_major_with_index():
        if c != ".":
            locations[c].add(yx)

    antinodes = set()
    for _, positions in locations.items():
        for (ya, xa), (yb, xb) in combinations(positions, 2):
            antinodes.add((ya, xa))
            antinodes.add((yb, xb))
            dy = yb - ya
            dx = xb - xa

            p1 = (ya - dy, xa - dx)
            while p1 in grid:
                antinodes.add(p1)
                p1 = (p1[0] - dy, p1[1] - dx)

            p2 = (yb + dy, xb + dx)
            while p2 in grid:
                antinodes.add(p2)
                p2 = (p2[0] + dy, p2[1] + dx)

    return len(antinodes)


print(part_one(data))
print(part_two(data))