# Finding the number of possible paths in a cave system (Python)

I've finally come up with working solution for day 12 of AdventOfCode, where you have to find the number of possible paths in a cave system following a given set of rules (described below). My solution seems a bit messy. I include an iterator class that employs recursive call of the same iterator. I'd be very grateful if you could review this code.

First I present the verbatim of the problem, then the input, then my solution.

With your submarine's subterranean subsystems subsisting suboptimally, the only way you're getting out of this cave anytime soon is by finding a path yourself. Not just a path - the only way to know if you've found the best path is to find all of them.

Fortunately, the sensors are still mostly working, and so you build a rough map of the remaining caves (your puzzle input). For example:

start-A
start-b
A-c
A-b
b-d
A-end
b-end


This is a list of how all of the caves are connected. You start in the cave named start, and your destination is the cave named end. An entry like b-d means that cave b is connected to cave d - that is, you can move between them.

So, the above cave system looks roughly like this:

    start
/   \
c--A-----b--d
\   /
end


Your goal is to find the number of distinct paths that start at start, end at end, > and don't visit small caves more than once. There are two types of caves: big caves (written in uppercase, like A) and small caves (written in lowercase, like b). It would be a waste of time to visit any small cave more than once, but big caves are large enough that it might be worth visiting them multiple times. So, all paths you find should visit small caves at most once, and can visit big caves any number of times.

Given these rules, there are 10 paths through this example cave system:

start,A,b,A,c,A,end
start,A,b,A,end
start,A,b,end
start,A,c,A,b,A,end
start,A,c,A,b,end
start,A,c,A,end
start,A,end
start,b,A,c,A,end
start,b,A,end
start,b,end


(Each line in the above list corresponds to a single path; the caves visited by that path are listed in the order they are visited and separated by commas.)

Note that in this cave system, cave d is never visited by any path: to do so, cave b would need to be visited twice (once on the way to cave d and a second time when returning from cave d), and since cave b is small, this is not allowed.

Here is a slightly larger example:

dc-end
HN-start
start-kj
dc-start
dc-HN
LN-dc
HN-end
kj-sa
kj-HN
kj-dc


The 19 paths through it are as follows:

start,HN,dc,HN,end
start,HN,dc,HN,kj,HN,end
start,HN,dc,end
start,HN,dc,kj,HN,end
start,HN,end
start,HN,kj,HN,dc,HN,end
start,HN,kj,HN,dc,end
start,HN,kj,HN,end
start,HN,kj,dc,HN,end
start,HN,kj,dc,end
start,dc,HN,end
start,dc,HN,kj,HN,end
start,dc,end
start,dc,kj,HN,end
start,kj,HN,dc,HN,end
start,kj,HN,dc,end
start,kj,HN,end
start,kj,dc,HN,end
start,kj,dc,end


Finally, this even larger example has 226 paths through it:

fs-end
he-DX
fs-he
start-DX
pj-DX
end-zg
zg-sl
zg-pj
pj-he
RW-he
fs-DX
pj-RW
zg-RW
start-pj
he-WI
zg-he
pj-fs
start-RW


How many paths through this cave system are there that visit small caves at most once?

Your puzzle answer was 4912.

--- Part Two ---

After reviewing the available paths, you realize you might have time to visit a single > small cave twice. Specifically, big caves can be visited any number of times, a single small cave can be visited at most twice, and the remaining small caves can be visited at most once. However, the caves named start and end can only be visited exactly once each: once you leave the start cave, you may not return to it, and once you reach the end cave, the path must end immediately.

Now, the 36 possible paths through the first example above are:

start,A,b,A,b,A,c,A,end
start,A,b,A,b,A,end
start,A,b,A,b,end
start,A,b,A,c,A,b,A,end
start,A,b,A,c,A,b,end
start,A,b,A,c,A,c,A,end
start,A,b,A,c,A,end
start,A,b,A,end
start,A,b,d,b,A,c,A,end
start,A,b,d,b,A,end
start,A,b,d,b,end
start,A,b,end
start,A,c,A,b,A,b,A,end
start,A,c,A,b,A,b,end
start,A,c,A,b,A,c,A,end
start,A,c,A,b,A,end
start,A,c,A,b,d,b,A,end
start,A,c,A,b,d,b,end
start,A,c,A,b,end
start,A,c,A,c,A,b,A,end
start,A,c,A,c,A,b,end
start,A,c,A,c,A,end
start,A,c,A,end start,A,end
start,b,A,b,A,c,A,end
start,b,A,b,A,end
start,b,A,b,end
start,b,A,c,A,b,A,end
start,b,A,c,A,b,end
start,b,A,c,A,c,A,end
start,b,A,c,A,end
start,b,A,end
start,b,d,b,A,c,A,end
start,b,d,b,A,end
start,b,d,b,end
start,b,end


The slightly larger example above now has 103 paths through it, and the even larger example now has 3509 paths through it.

Given these new rules, how many paths through this cave system are there?

Your puzzle answer was 150004.

The input:

YW-end
DK-la
la-XG
end-gy
zq-ci
XG-gz
TF-la
xm-la
gy-gz
ci-start
YW-ci
TF-zq
ci-DK
la-TS
zq-YW
gz-YW
zq-gz
end-gz
ci-TF
DK-zq
gy-YW
start-DK
gz-DK
zq-la
start-TF


The solution

import logging

# Debug header that enables logging if DEBUG = True. WARNING: A huge amount of data is being logged due to the sheer
# amount of computation, enabling DEBUG causes this program to take 10+ to execute and creates a ~500MB text file.
DEBUG = False

logger = logging.getLogger("new_service")
logger.setLevel(logging.FATAL)
if DEBUG:
logger.setLevel(logging.DEBUG)
log_format = logging.Formatter("%(message)s")

log_handler = logging.FileHandler("log12.txt", mode='w')
log_handler.setLevel(logging.DEBUG)
log_handler.setFormatter(log_format)

class CaveSystem:
"""
The class for keeping track of caves and their connections. Iterating through a CaveSystem instance results in
a list of paths where big caves can be visited unlimited amount of times and small caves can only be visited once.
Calling the same instance also returns an iterator, but if takes True as an argument, the number of paths are
extended by including those where a single small cave is explored twice.

Attributes:
caves : dict of caves in the instance, with cave names for keys and Cave objects for values.
"""

def __init__(self, puzzle):
"""
:param puzzle: list of str representations of cave connections
"""

# record every connection in a temporary list joint
joints = []
for joint in puzzle:
joints.append(joint.split('-'))

# extract the total amount of caves from joints and create a Cave object corresponding to every one of
# them. Register each cave in the self.caves attribute.
caves = set(flatten(joints))
self.caves = {}
for cave in caves:
self.caves[cave] = self.Cave(cave)

# Register every joint in each of the two corresponding Cave objects
for joint in joints:
self.caves[joint[0]].joints.append(joint[1])
self.caves[joint[1]].joints.append(joint[0])

def __call__(self, extend_one_small=False):
"""
Return a CaveIterator that iterates through every path that can be taken in the given cave system.
Same as self.__iter__() if extend_one_small is False, otherwise one single cave can be explore up to
2 times.
:param extend_one_small: bool, enables exploring one small cave up to two times in a single path.
Default False
:return: CaveIterator
"""
if extend_one_small:
return CaveIterator(self, [self.caves['start']], extend_one_small=True)
return self.__iter__()

def __iter__(self):
"""
Return a CaveIterator that iterates through every path that can be taken in the given cave system. Small
caves can only be explored once, while large caves can be explored any amount of times.
:return: CaveIterator
"""
# Every path begins at cave 'start'
return CaveIterator(self, [self.caves['start']])

class Cave:
"""
This nested class can only be used by the outer class and serves to store every cave within the same outer class

Attributes:
name : str the name of the cave. Capital letters designate large caves, and lowercase letters
designate small caves

big : bool, True if the cave is big, False otherwise

joints : list of str, listing every cave name that is connected to this one
"""

def __init__(self, name):
"""
:param name: str name of the cave. Capital letters designate the cave as being large, otherwise the cave
is considered small
"""
self.name = name
self.big = False
if self.name.lower() != self.name:
self.big = True

# Since only the name of the cave is supplied to the constructor, it is up to the outer class to fill this
# in
self.joints = []

class CaveIterator:
"""
A recursive iterator that goes deeper and deeper into recursion until it finds the end cave, which it then
returns. Then it returns the next unexplored path leading to the end cave. By default, small caves can only be
explored once. This can be changed by supplying the extend_one_small=True argument to the constructor, in
which case one small cave can be explored up to two times in a single path.

Attributes:
self.cave_system : CaveSystem that is being explored
self.entrance : list of Cave objects. Serves as the first, imutable part of every path return by this
instance of CaveIterator
self.extend_one_small : bool, tracks whether one small cave is allowed to be eplored up to two
times.
self.to_explore : list of Cave objects that that can be reached from the last cave in self.entrance,
that is, the current cave. This list is popped one by one throughout iteration.
"""

def __init__(self, cave_system, entrance, extend_one_small=False):
"""
:param cave_system: CaveSystem to be explored
:param entrance: list of Cave objects that serve as the first, immutable part of every apth generated by
the iterator
:param extend_one_small: bool, if True, one small cave can be explored twice in a path. Default is
False; every small cave can be explored no more than once.
"""

# Copy the initializing arguments
self.cave_system = cave_system
self.entrance = entrance
self.extend_one_small = extend_one_small

# Produce the self.to_explore attribute by creating the list of every Cave that is connected to the
# final one in self.entrance. The iterator will explore each of the caves in reverse order.
# If the 'end' cave is one of such cases, move it to the beginning of the list, making it the last cave to be
# explored. Note that Cave.joints is a list of str, thus the cave in the list expressions is a str
self.to_explore = [self.cave_system.caves[cave] for cave in self.entrance[-1].joints]
names_to_explore = caves_to_names(self.to_explore)
if 'end' in names_to_explore:
end = self.to_explore.pop(names_to_explore.index('end'))
self.to_explore.insert(0, end)

# Even when extend_one_small is toggled, the 'start' cave, which is technically a small cave, can still only
# be explored once, upon entrance, which should have happened by now, so it chould never be explored by the
# CaveIterator
if 'start' in names_to_explore:
del self.to_explore[names_to_explore.index('start')]

# If extend_one_small is not toggled, exclude every small cave from self.to_explore if it's already listed
# in self.entrance.
if not self.extend_one_small:
self._exclude_small_caves()
# Otherwise, see of there's a small cave that is already listed twice in self.entrance. Only then get rid of
# small caves in self.to_explore that are also listed in self.entrance
else:
for cave in self.entrance:
if not cave.big and self.entrance.count(cave) > 1:
self._exclude_small_caves()
logger.debug(f"{cave.name} triggers exclusion of small caves")
elif not cave.big:
logger.debug(f"{cave.name} is allowed twice")

# These two are used by self.__next__(), the first one keeping track of the currently explored cave, and
# the second one generating a nested CaveIterator for the cave currently being explored.
self._current_path = []
self._current_iterator = None

def __iter__(self):
"""
Return the actual iterator.
:return: iterator
"""
return self

def _exclude_small_caves(self):
"""
Exclude every small cave from self.to_explore that has already been listed in self.entrance.
:return: None
"""
for cave in self.entrance[:]:
if not cave.big and cave in self.to_explore:
self.to_explore.remove(cave)

def __next__(self):
"""
Produces the next path (list of str names in the order in which they are explored) for the current
CaveSystem that satisfies supplied conditions
:return:
"""

logger.debug(f"entrance: {caves_to_names(self.entrance)}, to explore: {caves_to_names(self.to_explore)}")

# If there're no caves left unexplored and the current cave has been fully explored, stop iteration
if not self.to_explore and not self._current_iterator:
raise StopIteration

# If the last cave left unexplored is the 'end' cave, append that cave to self.entrance and return it as the
# discovered path
if self.to_explore and self.to_explore[-1].name == 'end':
self.entrance.append(self.to_explore.pop())
return_path = caves_to_names(self.entrance)
return return_path
# If no cave is currently being explored, create an iterator for the last cave in self.to_explore
if not self._current_iterator:
self._current_path = self.entrance[:]
self._current_path.append(self.to_explore.pop())
self._current_iterator = CaveIterator(self.cave_system, self._current_path, self.extend_one_small)
# Try returning a new path going through self._current_path. If no such path can be found, reset
# self._current_iterator and try again by calling self.__next__() until it yields a path or stops iteration.
try:
return next(self._current_iterator)
except StopIteration:
self._current_iterator = None
return self.__next__()

def flatten(x):
"""
Flatten a list or a tuple
:param x: multidimensional list or tuple
:return: one-dimensional list of objects from the original x iterable
"""
flattened = []
for item in x:
if isinstance(item, list) or isinstance(item, tuple):
flattened.extend(flatten(item))
else:
flattened.append(item)
return flattened

def caves_to_names(lst):
"""
A simple function to produce a list of cave names (str) from a list of Cave objects
:param lst: a list of Cave objects
:return: a list of str cave names in the same order
"""
return [cave.name for cave in lst]

with open('input12.txt') as f:

# Split the data line by line and generate a cave system
puzzle = data.split('\n')
cave_system = CaveSystem(puzzle)

# This snippet of code lists every cave and its connections, for further debugging
if DEBUG:
for cave in cave_system.caves:
print(cave, cave_system.caves[cave].joints)

# Do the actual count
count = 0
for path in cave_system:
count += 1

extended_count = 0
for path in cave_system(extend_one_small=True):
logger.debug(f"{path}")
extended_count += 1

print(f"When every small caves is explored no more than once, the number of paths is {count},\n"
f"but when one small cave can be explored up to two times ina  single path, it's {extended_count}.")