80 lines
1.6 KiB
Python
80 lines
1.6 KiB
Python
import fileinput
|
|
|
|
def read_file():
|
|
return [line.strip() for line in fileinput.input()]
|
|
|
|
def make_graph(lines):
|
|
graph = {}
|
|
small_caves = set()
|
|
for line in lines:
|
|
a, b = line.split('-')
|
|
add_path(graph, a, b)
|
|
add_path(graph, b, a)
|
|
if a.islower():
|
|
small_caves.add(a)
|
|
if b.islower():
|
|
small_caves.add(b)
|
|
return graph, small_caves
|
|
|
|
|
|
def add_path(graph, a, b):
|
|
nodes = graph.get(a)
|
|
if nodes == None:
|
|
graph[a] = [b]
|
|
else:
|
|
nodes.append(b)
|
|
|
|
|
|
def find_paths(graph,
|
|
small_caves,
|
|
can_visit_twice,
|
|
current,
|
|
end,
|
|
path,
|
|
small_visited,
|
|
did_visit_twice=False):
|
|
|
|
if current in small_caves:
|
|
if current in small_visited:
|
|
if can_visit_twice and not did_visit_twice and current != 'start' and current != 'end':
|
|
did_visit_twice = True
|
|
else:
|
|
return None
|
|
small_visited.add(current)
|
|
|
|
path.append(current)
|
|
|
|
if current == end:
|
|
return [path]
|
|
|
|
paths = []
|
|
connected = graph[current]
|
|
for cave in connected:
|
|
subpaths = find_paths(graph, small_caves, can_visit_twice, cave, end,
|
|
path.copy(), small_visited.copy(), did_visit_twice)
|
|
if subpaths != None:
|
|
paths.extend(subpaths)
|
|
|
|
return paths
|
|
|
|
|
|
def print_graph(graph):
|
|
for key, value in graph.items():
|
|
for cave in value:
|
|
print(f"{key} -> {cave}")
|
|
|
|
|
|
def print_paths(paths):
|
|
for path in paths:
|
|
print(','.join(path))
|
|
|
|
|
|
input = read_file()
|
|
graph, small_caves = make_graph(input)
|
|
|
|
paths = find_paths(graph, small_caves, False, 'start', 'end', [], set())
|
|
print(f"Part 1: {len(paths)}")
|
|
|
|
paths = find_paths(graph, small_caves, True, 'start', 'end', [], set())
|
|
print(f"Part 2: {len(paths)}")
|
|
|