65 lines
1.5 KiB
Python
65 lines
1.5 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
import fileinput
|
||
|
|
from timeit import timeit
|
||
|
|
from typing import FrozenSet, List, Optional, Tuple, Union
|
||
|
|
from collections import defaultdict
|
||
|
|
|
||
|
|
lines = list(fileinput.input())
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
# parse
|
||
|
|
graph = defaultdict(list)
|
||
|
|
for line in lines:
|
||
|
|
fro, to = str(line).strip().split("-")
|
||
|
|
graph[fro].append(to)
|
||
|
|
graph[to].append(fro)
|
||
|
|
|
||
|
|
# part 1
|
||
|
|
stack: List[Tuple[str, FrozenSet[str]]] = [("start", frozenset())]
|
||
|
|
part1 = 0
|
||
|
|
while stack:
|
||
|
|
pos, visited = stack.pop()
|
||
|
|
if pos.islower() and pos in visited:
|
||
|
|
continue
|
||
|
|
visited = visited | {pos}
|
||
|
|
|
||
|
|
if pos == "end":
|
||
|
|
part1 += 1
|
||
|
|
continue
|
||
|
|
|
||
|
|
for next_pos in graph[pos]:
|
||
|
|
stack.append((next_pos, visited))
|
||
|
|
|
||
|
|
stack2: List[Tuple[str, FrozenSet[str], Optional[str]]] = [
|
||
|
|
("start", frozenset(), None)
|
||
|
|
]
|
||
|
|
part2 = 0
|
||
|
|
while stack2:
|
||
|
|
pos, visited, doubled = stack2.pop()
|
||
|
|
|
||
|
|
if pos.islower() and pos in visited:
|
||
|
|
if doubled is None:
|
||
|
|
stack2.append((pos, visited, pos))
|
||
|
|
continue
|
||
|
|
if doubled != pos:
|
||
|
|
continue
|
||
|
|
doubled = ""
|
||
|
|
|
||
|
|
visited = visited | {pos}
|
||
|
|
|
||
|
|
if pos == "end":
|
||
|
|
part2 += 1
|
||
|
|
continue
|
||
|
|
|
||
|
|
for next_pos in graph[pos]:
|
||
|
|
if next_pos != "start":
|
||
|
|
stack2.append((next_pos, visited, doubled))
|
||
|
|
|
||
|
|
return part1, part2
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
print(main())
|
||
|
|
print(f"{timeit(main, number=20) * 50:.3f} ms")
|