def parse(file): grid = [] for line in file: grid.append(line.strip()) return grid import itertools import string signal_chars = string.ascii_letters + string.digits def solve(file): grid = parse(file) ny = len(grid) nx = len(grid[0]) antennae = {} for y in range(ny): for x in range(nx): c = grid[y][x] if c in signal_chars: antennae.setdefault(c, []).append((x,y)) poles = set() for c, locations in antennae.items(): for (x0,y0), (x1,y1) in itertools.combinations(locations, 2): x = x1 + (x1-x0) y = y1 + (y1-y0) #print(c, (x0,y0),(x1,y1), "=>", (x,y)) if 0 <= x < nx and 0 <= y < ny: poles.add((x,y)) x = x0 + (x0-x1) y = y0 + (y0-y1) if 0 <= x < nx and 0 <= y < ny: poles.add((x,y)) #for y,line in enumerate(grid): # print("".join(c if (x,y) not in poles else '#' for (x,c) in enumerate(line))) print(len(poles)) poles = set() for c, locations in antennae.items(): for (x0,y0), (x1,y1) in itertools.combinations(locations, 2): for _ in 1,2: dx = (x1-x0) dy = (y1-y0) x = x1 y = y1 while 0 <= x < nx and 0 <= y < ny: #print(c, (x0,y0),(x1,y1), "=>", (x,y)) poles.add((x,y)) x += dx y += dy x0,y0, x1,y1 = x1,y1, x0,y0 #for y,line in enumerate(grid): # print("".join(c if (x,y) not in poles else '#' for (x,c) in enumerate(line))) print(len(poles)) solve(open("sample1.in")) solve(open("input"))