aoc-2023/d08/main.py

121 lines
2.8 KiB
Python
Raw Normal View History

2023-12-08 17:36:04 +00:00
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)