Advent of code 2024 day 4

def day4(s, part2=False):
  grid = np.array([list(line) for line in s.splitlines()])
  count = 0

  for _ in range(4):
    grid = np.rot90(grid)
    if not part2:
      count += ((grid[:-3] == 'X') & (grid[1:-2] == 'M') &
                (grid[2:-1] == 'A') & (grid[3:] == 'S')).sum()
      count += ((grid[:-3, :-3] == 'X') & (grid[1:-2, 1:-2] == 'M') &
                (grid[2:-1, 2:-1] == 'A') & (grid[3:, 3:] == 'S')).sum()
    else:
      count += ((grid[:-2, :-2] == 'M') & (grid[:-2, 2:] == 'M') &
                (grid[2:, :-2] == 'S') & (grid[2:, 2:] == 'S') &
                (grid[1:-1, 1:-1] == 'A')).sum()

  return count