import numpy as np verbose = False def load_map(filename="input"): """ Reads a map from a file without verifying its contents. """ with open(filename, "r") as fp: data = fp.read().splitlines() chararray = np.array(data, dtype=str)\ .view("U1")\ .reshape((len(data), -1)) return chararray def isValid(shape: tuple, index: tuple): for i in range(len(shape)): if not (0 <= index[i] < shape[i]): return False return True if __name__ == "__main__": m = load_map() frequencies = np.unique(m) antinodes = [] for f in frequencies: # Emptiness is killing me. if f == ".": continue # Find all antennas with the same frequency. antennas = np.array(np.where(m == f)).T # For each pair of antennas, find antinodes. for i in range(len(antennas)): for j in range(i + 1, len(antennas)): a = antennas[i] b = antennas[j] c = a + (a - b) d = b + (b - a) if isValid(m.shape, c): antinodes.append(c) if isValid(m.shape, d): antinodes.append(d) print(len(np.unique(antinodes,axis=0)))