87 lines
2.6 KiB
Python
87 lines
2.6 KiB
Python
import numpy as np
|
|
from pathlib import Path
|
|
|
|
|
|
def pairs_from_list(array):
|
|
if len(array) == 2:
|
|
return [tuple(array)]
|
|
return [(array[0], k) for k in array[1:]] + pairs_from_list(array[1:])
|
|
|
|
|
|
class AntennaMap:
|
|
def __init__(self, filepath):
|
|
with open(filepath, "r") as filein:
|
|
self.map = np.array(
|
|
[[ord(c) for c in line.rstrip()] for line in filein.readlines()]
|
|
)
|
|
self.xlen, self.ylen = self.map.shape
|
|
unifrq = np.unique(self.map)
|
|
unique_freqs = np.delete(unifrq, np.argwhere(unifrq == ord(".")))
|
|
self.nodes = set()
|
|
self.antennae = {k: np.argwhere(self.map == k) for k in unique_freqs}
|
|
|
|
def __str__(self):
|
|
map = np.copy(self.map)
|
|
for i, j in self.nodes:
|
|
map[i, j] = ord("#") if map[i, j] == ord(".") else map[i, j]
|
|
return "\n".join(["".join([chr(c) for c in line]) for line in map])
|
|
|
|
def _on_map(self, coords):
|
|
x, y = coords
|
|
return (0 <= x < self.xlen) and (0 <= y < self.ylen)
|
|
|
|
def find_antinode_pairs(self, freq):
|
|
coords = self.antennae[freq]
|
|
pairs = pairs_from_list(coords)
|
|
for x1, x2 in pairs:
|
|
u = np.subtract(x2, x1)
|
|
x3 = x2 + u
|
|
x0 = x1 - u
|
|
if self._on_map(x0):
|
|
self.nodes.add(tuple(x0))
|
|
if self._on_map(x3):
|
|
self.nodes.add(tuple(x3))
|
|
|
|
def _border_distance_2d(self, coords, vect):
|
|
result = [
|
|
self._border_distance_1d(x, u, mx)
|
|
for x, u, mx in zip(coords, vect, self.map.shape)
|
|
]
|
|
return np.min(result)
|
|
|
|
def _border_distance_1d(self, x, u, mx):
|
|
xb = 0 if u < 0 else mx - 1
|
|
length = xb - x
|
|
return length // u
|
|
|
|
def find_antinodes(self, freq):
|
|
coords = self.antennae[freq]
|
|
pairs = pairs_from_list(coords)
|
|
for x1, x2 in pairs:
|
|
u = np.subtract(x2, x1)
|
|
k0 = self._border_distance_2d(x2, -u)
|
|
k1 = self._border_distance_2d(x2, u)
|
|
for k in range(-k0, k1 + 1):
|
|
self.nodes.add(tuple(x2 + k * u))
|
|
|
|
def scan_prime(self):
|
|
for freq in self.antennae:
|
|
self.find_antinode_pairs(freq)
|
|
|
|
def scan_harmonics(self):
|
|
for freq in self.antennae:
|
|
self.find_antinodes(freq)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
filepath = Path("./example")
|
|
filepath = Path("./input")
|
|
|
|
map = AntennaMap(filepath)
|
|
map.scan_prime()
|
|
print(map)
|
|
print(len(map.nodes))
|
|
|
|
map.scan_harmonics()
|
|
print(map)
|
|
print(len(map.nodes))
|