Advent of code 2022 day 17

def day17(s, *, part2=False, num=10**12, rocks_in_part2=8_000):
  rocks = [np.array(rock, bool)[::-1] for rock in [
    [[1, 1, 1, 1]],
    [[0, 1, 0], [1, 1, 1], [0, 1, 0]],
    [[0, 0, 1], [0, 0, 1], [1, 1, 1]],
    [[1], [1], [1], [1]],
    [[1, 1], [1, 1]]
  ]]
  jet_iterator = itertools.cycle(s)
  num_rocks = rocks_in_part2 if part2 else 2022
  grid = np.full((num_rocks * 4, 7), False)  # (y=0 is at bottom.)
  heights = [0]

  for rock in itertools.islice(itertools.cycle(rocks), num_rocks):
    ry, rx = rock.shape
    y, x = heights[-1] + 3, 2
    while True:
      dx = {'<': -1, '>': 1}[next(jet_iterator)]
      if (x + dx >= 0 and x + dx + rx <= 7 and
          not (grid[y:y + ry, x + dx: x + dx + rx] & rock).any()):
        x += dx
      if y > 0 and not (grid[y - 1:y - 1 + ry, x:x + rx] & rock).any():
        y -= 1
      else:
        grid[y:y + ry, x:x + rx] |= rock
        heights.append(max(heights[-1], y + ry))
        break

  if not part2:
    return heights[-1]

  # Find a repeating sequence of height deltas, starting from the end.
  deltas = np.diff(heights)
  tail = deltas[-rocks_in_part2 // 3:]
  # If next line fails, increase `rocks_in_part2`.
  i0, i1, *_ = [i for i in range(len(deltas) - len(tail) + 1)
                if all(deltas[i:i + len(tail)] == tail)]
  h0, h1 = heights[i0], heights[i1]
  div, mod = divmod(num - i0, i1 - i0)
  return h0 + div * (h1 - h0) + heights[i0 + mod] - h0