88 lines
2.4 KiB
Python
Executable File
88 lines
2.4 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
import sys
|
|
from collections import defaultdict
|
|
from copy import deepcopy
|
|
from typing import DefaultDict, MutableMapping, Set
|
|
import pprint
|
|
import string
|
|
|
|
|
|
Vertices = DefaultDict[str, Set[str]]
|
|
SMALL = set(string.ascii_lowercase)
|
|
BIG = set(string.ascii_uppercase)
|
|
|
|
|
|
def is_small(cave: str):
|
|
s = set(cave)
|
|
return (s & SMALL) == s
|
|
|
|
|
|
def part1(vertices: Vertices):
|
|
small = {cave for cave in vertices.keys() if is_small(cave)}
|
|
big = set(vertices.keys()) - small
|
|
assert len(big & small) == 0
|
|
|
|
def find_paths(current: str, visited: Set[str] = set(), path: str = "") -> Set[str]:
|
|
visited = deepcopy(visited)
|
|
visited |= {current}
|
|
paths = set()
|
|
for vertex in vertices[current]:
|
|
if vertex == "end":
|
|
paths |= {path}
|
|
elif vertex in visited and vertex in small:
|
|
continue
|
|
else:
|
|
paths |= find_paths(vertex, visited, path + vertex)
|
|
return paths
|
|
|
|
paths = find_paths("start")
|
|
print(len(paths))
|
|
|
|
|
|
def part2(vertices: Vertices):
|
|
small = {cave for cave in vertices.keys() if is_small(cave)}
|
|
big = set(vertices.keys()) - small
|
|
assert len(big & small) == 0
|
|
|
|
def find_paths(
|
|
allowed: str,
|
|
current: str,
|
|
visited: MutableMapping[str, int] = defaultdict(lambda: 0),
|
|
path: str = "",
|
|
) -> Set[str]:
|
|
if current == allowed:
|
|
if visited[current] >= 2:
|
|
# This is a small cave and we've already visited it twice
|
|
return set()
|
|
elif current in small and visited[current]:
|
|
# Visited already
|
|
return set()
|
|
visited = deepcopy(visited)
|
|
visited[current] += 1
|
|
paths = set()
|
|
for vertex in vertices[current]:
|
|
if vertex == "end":
|
|
paths |= {path}
|
|
else:
|
|
paths |= find_paths(allowed, vertex, visited, path + vertex)
|
|
return paths
|
|
|
|
paths = set()
|
|
for allowed in small - {"start", "end"}:
|
|
paths |= find_paths(allowed, "start")
|
|
print(len(paths))
|
|
|
|
|
|
# Doing this comprehension so Mypy is happy
|
|
edges = [(a, b) for a, b in map(lambda l: l.strip().split("-"), sys.stdin)]
|
|
vertices: Vertices = defaultdict(set)
|
|
for a, b in edges:
|
|
vertices[a] |= {b}
|
|
vertices[b] |= {a}
|
|
|
|
|
|
print("Part 1")
|
|
part1(vertices)
|
|
print("Part 2")
|
|
part2(vertices)
|