121 lines
2.8 KiB
Python
121 lines
2.8 KiB
Python
from pathlib import Path
|
|
from typing import NamedTuple
|
|
from itertools import cycle
|
|
import math
|
|
|
|
|
|
class Node(NamedTuple):
|
|
id: str
|
|
left: str
|
|
right: str
|
|
|
|
|
|
def parse(input: Path) -> tuple[str, dict[str, Node]]:
|
|
with input.open() as f:
|
|
|
|
steps = f.readline().strip()
|
|
nodes: dict[str, Node] = {}
|
|
|
|
_ = f.readline()
|
|
|
|
for line in f:
|
|
node = Node(
|
|
line[0:3],
|
|
line[7:10],
|
|
line[12:15],
|
|
)
|
|
nodes[node.id] = node
|
|
|
|
return steps, nodes
|
|
|
|
|
|
def part2(input: Path) -> int:
|
|
steps, nodes = parse(input)
|
|
|
|
start_nodes: list[str] = [
|
|
node for node in nodes.keys() if node.endswith("A")
|
|
]
|
|
|
|
periods: dict[str, int] = {}
|
|
for start_node in start_nodes:
|
|
current = start_node
|
|
num_steps = 0
|
|
for step in cycle(steps):
|
|
if current.endswith("Z"):
|
|
print(f"Found a node for {start_node} at {current} in {num_steps}")
|
|
periods[start_node] = num_steps
|
|
# This assumes there is only one possible end and that it begins a
|
|
# cycle after that. Apparently true for this data set.
|
|
break
|
|
|
|
num_steps += 1
|
|
|
|
if step == "L":
|
|
current = nodes[current].left
|
|
elif step == "R":
|
|
current = nodes[current].right
|
|
else:
|
|
raise ValueError(f"Unknown step {step}")
|
|
|
|
return math.lcm(*[p for p in periods.values()])
|
|
|
|
|
|
def part2_brute(input: Path) -> int:
|
|
steps, nodes = parse(input)
|
|
|
|
current: list[str] = [
|
|
node for node in nodes.keys() if node.endswith("A")
|
|
]
|
|
|
|
num_steps = 0
|
|
for step in steps:
|
|
if all(node.endswith("Z") for node in current):
|
|
break
|
|
|
|
num_steps += 1
|
|
|
|
for i, node in enumerate(current):
|
|
# if node.endswith("Z"):
|
|
# print(f"Node {initial[i]} reached {node} in {num_steps-1} steps")
|
|
if step == "L":
|
|
current[i] = nodes[node].left
|
|
elif step == "R":
|
|
current[i] = nodes[node].right
|
|
else:
|
|
raise ValueError(f"Unknown step {step}")
|
|
|
|
return num_steps
|
|
|
|
|
|
def part1(input: Path) -> int:
|
|
steps_s, nodes = parse(input)
|
|
steps = cycle(steps_s)
|
|
|
|
current = "AAA"
|
|
num_steps = 0
|
|
for step in steps:
|
|
if current == "ZZZ":
|
|
break
|
|
|
|
num_steps += 1
|
|
|
|
if step == "L":
|
|
current = nodes[current].left
|
|
elif step == "R":
|
|
current = nodes[current].right
|
|
else:
|
|
raise ValueError(f"Unknown step {step}")
|
|
|
|
return num_steps
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# input = Path("sample.txt")
|
|
input = Path("input.txt")
|
|
|
|
result1 = part1(input)
|
|
print("part1", result1)
|
|
|
|
result2 = part2(input)
|
|
print("part2", result2)
|