Advent of code 2023 day 23

def day23_graph(s, part2):
  grid = np.array([list(line) for line in s.splitlines()])
  start_yx, target_yx = (0, 1), (grid.shape[0] - 1, grid.shape[1] - 2)

  def neighbors(yx):
    y, x = yx
    results = []
    for dy, dx, ch in ((0, 1, '<'), (0, -1, '>'), (1, 0, '^'), (-1, 0, 'v')):
      if part2 or grid[yx] != ch:
        y2, x2 = y + dy, x + dx
        if 0 <= y2 < len(grid) and grid[y2, x2] != '#':
          results.append((y2, x2))
    return results

  nodes = {yx for yx, ch in np.ndenumerate(grid) if ch != '#' and len(neighbors(yx)) > 2}
  nodes |= {start_yx, target_yx}

  def get_paths(yx):
    paths = {}
    for yx2 in neighbors(yx):
      path = [yx]
      while yx2 not in nodes:
        path.append(yx2)
        candidates = [yx3 for yx3 in neighbors(yx2) if yx3 != path[-2]]
        (yx2,) = candidates if candidates else (yx,)  # We use yx to signal no legal path.
      if yx2 != yx:
        paths[yx2] = path[1:] + [yx2]
    return paths

  graph_yx = {yx: get_paths(yx) for yx in nodes}
  ordered = sorted(graph_yx)
  graph = [
      [(ordered.index(yx2), len(path)) for yx2, path in graph_yx[yx].items()]
      for i, yx in enumerate(ordered)
  ]
  start, target = ordered.index(start_yx), ordered.index(target_yx)
  return graph, start, target


@numba.njit
def day23_length_of_longest_path(graph_array, start, target):
  stack = [(start, 0, {start})]
  max_length = 0

  while stack:
    node, length, visited = stack.pop()
    if node == target:
      max_length = max(max_length, length)
      continue
    for node2, dist in graph_array[node]:
      if node2 >= 0 and node2 not in visited:
        stack.append((node2, length + dist, visited | {node2}))

  return max_length


def day23(s, *, part2=False):
  graph, start, target = day23_graph(s, part2)
  graph_array = np.full((len(graph), 4, 2), -1)  # np.int16 is no faster.
  for index, index2_paths in enumerate(graph):
    if index2_paths:
      graph_array[index, : len(index2_paths)] = index2_paths
  return day23_length_of_longest_path(graph_array, start, target)