Advent of code 2020 day 20

def process(s, part2=False):

  def to_grid(s):
    return np.array(list(map(list, s.strip('\n').splitlines())))

  tiles = {int(t[5:9]): to_grid(t[11:]) for t in s.strip('\n').split('\n\n')}
  n = int(len(tiles)**0.5)
  rotations = tuple(range(8))  # 4 proper rotations * 2 flips

  def rotate(tile, rotation):
    return np.rot90(tile if rotation < 4 else tile[::-1], rotation % 4)

  edge_count = collections.Counter(
      tuple(rotate(tile, rotation)[0])  # extract top edge
      for tile in tiles.values() for rotation in rotations
  )

  def is_corner(tile):
    # A corner tile has 2 unique edges, which appear 4 times after flips.
    return sum(edge_count[tuple(rotate(tile, rotation)[0])] == 1
               for rotation in rotations) == 4

  corners = [index for index, tile in tiles.items() if is_corner(tile)]
  assert len(corners) == 4
  if not part2:
    return np.prod(corners)

  index = corners[0]
  tile = tiles[index]
  rotation, _ = (rotation for rotation in rotations  # 2 solutions due to flip
                 if (edge_count[tuple(rotate(tile, rotation)[0])] == 1 and
                     edge_count[tuple(rotate(tile, rotation)[:, 0])] == 1))
  layout = np.empty((n, n), dtype=object)
  layout[0, 0] = (index, rotation)

  def find(not_index, subview, desired):
    "Returns a tile index and rotation such that its subview matches desired."
    for index in tiles:
      if index != not_index:
        for rotation in rotations:
          if np.all(rotate(tiles[index], rotation)[subview] == desired):
            return index, rotation
    raise AssertionError

  for y, x in np.ndindex(n, n):
    if x > 0:
      left_index, left_rotation = layout[y, x - 1]
      desired_left = rotate(tiles[left_index], left_rotation)[:, -1]
      layout[y, x] = find(left_index, np.s_[:, 0], desired_left)
    elif y > 0:
      top_index, top_rotation = layout[y - 1, x]
      desired_top = rotate(tiles[top_index], top_rotation)[-1, :]
      layout[y, x] = find(top_index, np.s_[0, :], desired_top)

  def block(y, x):
    index, rotation = layout[y, x]
    return rotate(tiles[index], rotation)[1:-1, 1:-1]

  grid = np.block([[block(y, x) for x in range(n)] for y in range(n)])

  pattern = '                  # #    ##    ##    ### #  #  #  #  #  #   '
  pattern = np.array(list(pattern)).reshape(3, -1) == '#'

  for rotation in rotations:
    gridview = rotate(grid, rotation)
    for y, x in np.ndindex(*np.array(gridview.shape) - pattern.shape):
      subgrid = gridview[y:y + pattern.shape[0], x:x + pattern.shape[1]]
      if np.all(~pattern | (subgrid != '.')):
        subgrid[pattern] = 'O'

  return np.count_nonzero(grid == '#')