Advent of code 2024 day 24 part 2

def day24_parse_vals_and_expressions(s):
  p1, p2 = s.split('\n\n')
  vals = {k: int(v) for line in p1.splitlines() for k, v in (line.split(': '),)}

  @dataclasses.dataclass
  class Expression:
    op: str
    srcs: tuple[str, str]

  expressions = {
      dst: Expression(op, (src1, src2))
      for line in p2.splitlines()
      for src1, op, src2, dst in (re.findall(r'\w+', line),)
  }

  return vals, expressions


def day24_part2(s, *, correct_to=1000, return_pairs=False):
  _, expressions = day24_parse_vals_and_expressions(s)
  n_z = sum(1 for dst in expressions if dst[0] == 'z')

  def orient(expr):  # Heuristically assign a canonical order for two expression operands.
    expr2 = expressions.get(expr.srcs[1], None)
    if expr.srcs[1][0] == 'x' or (expr2 and (expr2.op == 'XOR' or expr2.srcs[0][0] in 'xy')):
      expr.srcs = expr.srcs[::-1]

  def tree_is_correct(i):
    e = expressions[f'z{i:02}']
    if e.op == 'XOR' and all(n in expressions for n in e.srcs):
      e1, e2 = expressions[e.srcs[0]], expressions[e.srcs[1]]
      if e1.op == 'XOR' and e1.srcs[0][0] == 'x':
        if e2.op == 'OR' and all(n in expressions for n in e2.srcs):
          e21, e22 = expressions[e2.srcs[0]], expressions[e2.srcs[1]]
          if e21.op == 'AND' and e21.srcs[0][0] == 'x':
            last_e = expressions[f'z{i - 1:02}']
            if i <= 2 or e22.op == 'AND' and set(e22.srcs) == set(last_e.srcs):
              return True
    return False

  for expr in expressions.values():
    orient(expr)
  known_correct = set[str]()
  pairs = []

  for i in range(n_z):

    def local_nodes(i: int) -> Iterator[str]:
      seen = known_correct | set(expressions[f'z{i - 1:02}'].srcs)

      def recurse(dst: str) -> Iterator[str]:
        seen.add(dst)
        yield dst
        expr = expressions[dst]
        for n in expr.srcs:
          if n in expressions and n not in seen:
            yield from recurse(n)

      yield from recurse(f'z{i:02}')

    def locally_orient(i: int) -> None:
      for dst in local_nodes(i):
        orient(expressions[dst])

    def swap(dst1: str, dst2: str) -> None:
      expressions[dst1], expressions[dst2] = expressions[dst2], expressions[dst1]
      locally_orient(i)
      locally_orient(i + 1)

    if 2 <= i < n_z - 2:

      if i <= correct_to and not tree_is_correct(i):
        # print(f'Problem at z{i:02}')
        # We consider exchanging each pair within a set of local nodes, specifically those
        # contributing to the outputs i and i + 1, as this seems to be the puzzle pattern.
        dsts = set(local_nodes(i)) | set(local_nodes(i + 1))
        for dst1, dst2 in itertools.combinations(dsts, 2):
          swap(dst1, dst2)
          if tree_is_correct(i) and tree_is_correct(i + 1):
            # print(f'At {i=}, found a correcting pair {dst1}, {dst2}.')
            pairs.append((dst1, dst2))
            break
          swap(dst1, dst2)  # Swap back since it didn't solve it.
        else:
          raise AssertionError(f'Could not find repair for z{i:02}')

      known_correct.update(local_nodes(i))

  return pairs if return_pairs else ','.join(sorted(itertools.chain.from_iterable(pairs)))