
# Python imports
from collections import defaultdict
from copy import copy
from pathlib import Path
from typing import Callable, Dict, Generator, Iterable, List, Tuple

import networkx as nx

# Paths to data
testpath1 = Path("day12_test.txt")
testpath2 = Path("day12_test2.txt")
testpath3 = Path("day12_test3.txt")
datapath = Path("day12_data.txt")


def load_input(fpath: Path) -> List[Tuple[str, str]]:
    """Return graph as a list of edges
    
    :param fpath:  Path to data file
    """
    with fpath.open("r") as ifh:
        return [tuple(_.strip().split("-")) for _ in ifh.readlines() if len(_.strip())]
    
def graph_from_edges(edges: List[Tuple[str, str]]) -> nx.Graph:
    """Returns a graph corresponding to the edgelist
    
    :param edges:  list of edges in puzzle data
    """
    gph = nx.Graph()  # empty Graph
    
    # Add edges
    for start, end in edges:
        gph.add_edge(start, end)
    
    return gph


def find_paths(gph: nx.Graph) -> List[List[str]]:
    """Returns all paths from start to end that never touch a smallcave twice
    
    :param gph:  puzzle graph
    """
    paths = [(["start"], ["start",])]  # initialise paths
    fullpaths = []  # will hold paths that reach the end

    # Iterate until all test paths are discarded, or found to be
    # valid paths from start to end
    while len(paths):
        newpaths = []  # List of valid paths found in this iteration
        
        for smallcaves, path in paths:  # Check each active path
            # Iterate over neighbour node labels at end of path, if not
            # an already visited small cave
            for nbr in [str(_) for _ in gph[path[-1]] if _ not in smallcaves]:
                if nbr == "end":  # neighbour is end
                    # Extend path to end and add to fullpaths
                    fullpaths.append(path[:] + [nbr])
                elif nbr.lower() == nbr:  # neighbour is small cave
                    # Extend list of small caves, extend path, and add
                    # to list of paths for next iteration
                    newpaths.append((smallcaves[:] + [nbr], path[:] + [nbr]))
                else:
                    # Do not change list of small caves, extend path,
                    # and add to list of paths for next iteration
                    newpaths.append((smallcaves[:], path[:] + [nbr]))
                    
        # Update new paths for next iteration
        paths = newpaths[:]
        
    return fullpaths


edges = load_input(testpath1)
gph = graph_from_edges(edges)
len(find_paths(gph))

10


edges = load_input(testpath2)
gph = graph_from_edges(edges)
len(find_paths(gph))

19


edges = load_input(testpath3)
gph = graph_from_edges(edges)
len(find_paths(gph))

226


edges = load_input(datapath)
gph = graph_from_edges(edges)
len(find_paths(gph))

4413


def find_paths_with_one_small_revisit(gph: nx.Graph) -> List[List[str]]:
    """Returns all paths from start to end that never touch a smallcave twice
    
    :param gph:  puzzle graph
    """
    paths = [(defaultdict(int), ["start",])]  # initialise paths
    fullpaths = []  # will hold paths that reach the end

    # Iterate until all test paths are discarded, or found to be
    # valid paths from start to end
    while len(paths):
        newpaths = []  # list of valid paths found in this iteration
        maxvisits = 0  # initialise max visit count for small caves
        
        for smallcaves, path in paths:  # Check each active path
            # Do we need to update maxvisits?
            if maxvisits < 3 and len(smallcaves):
                maxvisits = max(smallcaves.values())
                
            # Iterate over neighbour node labels at end of path, if not
            # a small cave that has already been visited twice
            for nbr in [str(_) for _ in gph[path[-1]] if smallcaves[_] < 2]:
                if nbr == "end":  # neighbour is end
                    # Extend path to end and add to fullpaths
                    fullpaths.append(path[:] + [nbr])
                elif nbr.upper() == nbr:  # neigbour is large cave
                    # Do not change dict of small caves, extend path,
                    # and add to list of paths for next iteration
                    newpaths.append((copy(smallcaves), path[:] + [nbr]))
                elif nbr == "start":  # do not revisit start
                    pass
                else:  # neighbour is small cave
                    # Extend dict of small caves, extend path, and add
                    # to list of paths for next iteration
                    newcaves = copy(smallcaves)
                    if maxvisits < 2:  # No cave has yet been visited twice
                        newcaves[nbr] += 1                        
                        newpaths.append((newcaves, path[:] + [nbr]))
                    elif newcaves[nbr] == 0:  # This cave has not been visited
                        newcaves[nbr] += 1                        
                        newpaths.append((newcaves, path[:] + [nbr]))
                    else:  # Cave has been visited, and max revisits already reached
                        pass                        
                    
        # Update new paths for next iteration
        paths = newpaths[:]
        
    return fullpaths


edges = load_input(testpath1)
gph = graph_from_edges(edges)
len(find_paths_with_one_small_revisit(gph))

36


edges = load_input(testpath2)
gph = graph_from_edges(edges)
len(find_paths_with_one_small_revisit(gph))

103


edges = load_input(testpath3)
gph = graph_from_edges(edges)
len(find_paths_with_one_small_revisit(gph))

3509


edges = load_input(datapath)
gph = graph_from_edges(edges)
len(find_paths_with_one_small_revisit(gph))

118803


%%timeit
edges = load_input(datapath)
gph = graph_from_edges(edges)
len(find_paths_with_one_small_revisit(gph))

4.08 s ± 153 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Day 12: Passage Pathing¶

Puzzle 2:¶