1  from lib import *
 2  
 3  input = read_input(2021, 20)
 4  
 5  lines = input.splitlines()
 6  
 7  
 8  def solve(n):
 9      algo = [c == "#" for c in lines[0]]
10      grid = {(j, i) for i, line in enumerate(lines[2:]) for j, x in enumerate(line) if x == "#"}
11      inf = False
12      for _ in range(n):
13          new_inf = inf != algo[0]
14          candidates = {(p, q) for x, y in grid for p, q in get_neighbors(x, y, diag=True, include_self=True)}
15          new_grid = set()
16          for x, y in candidates:
17              idx = 0
18  
19              for p, q in sorted(get_neighbors(x, y, diag=True, include_self=True), key=lambda a: a[::-1]):
20                  idx <<= 1
21                  idx |= ((p, q) in grid) != inf
22  
23              if algo[idx] != new_inf:
24                  new_grid.add((x, y))
25  
26          grid = new_grid
27          inf = new_inf
28  
29      return len(grid)
30  
31  
32  print(solve(2))
33  print(solve(50))