aoc2024/d08/run.py

36 lines
938 B
Python

import sys, os
from collections import defaultdict as dd
from itertools import combinations as comb
from grid import *
R = list(range(50)) if len(sys.argv)>1 and sys.argv[1]=="2" else [1]
G = [list(l) for l in sys.stdin.read().strip().split("\n")]
# dict of antenna's position grouped by freq
P = dd(list)
W, H = grid_geom(G)
for y in range(H): # rows
for x in range(W): # cols
v = get_in_grid(G, x, y)
if v != '.':
P[v].append((x, y))
A = set() # antinodes
for _,v in P.items():
for a,b in comb(v,2): # for each pair of antenna resonating
x1, y1 = a
x2, y2 = b
dx = x2 - x1
dy = y2 - y1
for i in R:
x3 = x1 - i*dx
y3 = y1 - i*dy
x4 = x2 + i*dx
y4 = y2 + i*dy
if in_grid(G, x3, y3):
A.add((x3,y3))
if in_grid(G, x4, y4):
A.add((x4,y4))
print(len(A))