Advent of code 2023 day 8 general

def _extended_gcd(a: int, b: int) -> tuple[int, int, int]:
  """Finds the greatest common divisor using the extended Euclidean algorithm.

  Returns:
    (gcd(a, b), x, y) with the property that a * x + b * y = gcd(a, b).
  """
  prev_x, x = 1, 0
  prev_y, y = 0, 1
  while b:
    q = a // b
    x, prev_x = prev_x - q * x, x
    y, prev_y = prev_y - q * y, y
    a, b = b, a % b
  x, y = prev_x, prev_y
  return a, x, y


def _solve_modulo_congruences(moduli, remainders):
  """Returns `x` satisfying `x % moduli[i] == remainders[i]`; handles non-coprime moduli."""
  def merge(mr1, mr2):
    (m, a), (n, b) = mr1, mr2
    gcd, u, v = _extended_gcd(m, n)
    if 0:  # Simpler algorithm that assumes the moduli are coprime.
      return m * n, (a * v * n + b * u * m) % (m * n)
    else:  # General algorithm; see https://math.stackexchange.com/a/1644698.
      lamb = (a - b) // gcd
      sigma = a - m * u * lamb
      assert sigma == b + n * lamb * v
      lcm = math.lcm(m, n)
      return lcm, sigma % lcm

  return functools.reduce(merge, zip(moduli, remainders))[1]


def day8(s, *, part2=False):
  lines = s.splitlines()
  moves = lines[0]
  routes = {}
  for line in lines[2:]:
    src, dsts = line.split(' = ')
    routes[src] = dsts[1:-1].split(', ')

  def advance(node, move):
    return routes[node][{'L': 0, 'R': 1}[move]]

  if not part2:
    node = 'AAA'
    for index, move in enumerate(itertools.cycle(moves)):
      if node == 'ZZZ':
        return index
      node = advance(node, move)

  # Run an initial simulation long enough that all nodes should be cycling.
  nodes = {node for node in routes if node.endswith('A')}
  num = 1_000
  for index, move in enumerate(itertools.islice(itertools.cycle(moves), num)):
    if all(node.endswith('Z') for node in nodes):
      return index
    nodes = {advance(node, move) for node in nodes}

  # Identify the periods and phases of the cycles.
  moduli, remainders = [], []
  for node in nodes:
    last_seen = {}
    for index, move in enumerate(itertools.islice(itertools.cycle(moves), num, None)):
      if node.endswith('Z'):
        if node in last_seen:
          moduli.append(index - last_seen[node])
          break
        last_seen[node] = index
      node = advance(node, move)
    # Fortunately, the problem is special in that each cycle contains a single solution.
    assert len(last_seen) == 1
    remainders.append(last_seen[node])

  index = num + _solve_modulo_congruences(moduli, remainders)
  if 1:  # A particular property of this particular input:
    # The remainders happen to align such that the solution lies exactly at math.lcm(*moduli)!
    assert all(remainder + num == modulo for modulo, remainder in zip(moduli, remainders))
    assert index == math.lcm(*moduli)
  return index