Advent of code 2024 day 18 part2 fastest Union-Find

# import hhoppe_tools as hh  # For hh.UnionFind[tuple[int, int]].

# See https://en.wikipedia.org/wiki/Disjoint-set_data_structure.
class UnionFind:

  def __init__(self):
    self._rep = {}

  def union(self, a, b):
    rep_a, rep_b = self.find(a), self.find(b)
    self._rep[rep_b] = rep_a

  def same(self, a, b):
    return self.find(a) == self.find(b)

  def find(self, a):
    if a not in self._rep:
      return a
    parents = []
    while True:
      parent = self._rep.setdefault(a, a)
      if parent == a:
        break
      parents.append(a)
      a = parent
    for p in parents:
      self._rep[p] = a
    return a


# We detect if there an 8-connected path between anywhere on the N,E boundaries and anywhere on
# the S,W boundaries.  If there is such a connected path, then we know that it prevents a
# 4-connected path from (0, 0) to (size - 1, size - 1).
# We iteratively consider each block yx, applying union_find.union(yx, yx2) with the special
# labels nw_boundary and se_boundary if yx is boundary-adjacent and with all its 8-connected
# neighboring blocks.
# After adding each block, we see if the two boundary components are now in the same connected set.
def day18_part2(s, *, size=71):
  blocks = (map(int, row.split(',')) for row in s.splitlines())
  union_find = UnionFind()
  neighbors = set(itertools.product((-1, 0, 1), repeat=2)) - {(0, 0)}
  seen = set()
  nw_boundary, se_boundary = (-1, size), (size, -1)  # Arbitrary labels.
  for x, y in blocks:
    seen.add((y, x))
    for dy, dx in neighbors:
      y2, x2 = y + dy, x + dx
      if y2 < 0 or x2 >= size:
        union_find.union((y, x), nw_boundary)
      elif x2 < 0 or y2 >= size:
        union_find.union((y, x), se_boundary)
      elif (y2, x2) in seen:
        union_find.union((y, x), (y2, x2))
    if union_find.same(nw_boundary, se_boundary):
      return f'{x},{y}'