from lib import get_data from lib import Grid2D data = get_data(__file__) locks = [] keys = [] for o in data.split("\n\n"): g = Grid2D(o) if all(g[(0, c)] == "#" for c in range(g.n_cols)): keys.append(g) elif all(g[g.n_rows - 1, c] == "#" for c in range(g.n_cols)): locks.append(g) else: assert False t = 0 for key in keys: for lock in locks: k = key.find("#") l = lock.find("#") if len(k) + len(l) == len(set(k) | set(l)): t += 1 print(t)