from lib import get_data from lib import Grid2D from itertools import combinations from collections import defaultdict data = get_data(__file__) g = Grid2D(data) d = defaultdict(list) for row in range(g.n_rows): for col in range(g.n_cols): p = (row, col) if g[p] != ".": d[g[p]].append((row, col)) aa = set() bb = set() for coords in d.values(): for c, d in combinations(coords, 2): bb.add(c) bb.add(d) for a, b in [(c, d), (d, c)]: d = b[0] - a[0], b[1] - a[1] a = b[0] + d[0], b[1] + d[1] if g.contains(a): aa.add(a) a = b[0] + d[0], b[1] + d[1] while g.contains(a): bb.add(a) a = a[0] + d[0], a[1] + d[1] print(len(aa)) print(len(bb))