from collections import defaultdict

# build the adjacency graph from the input data
def build_graph(data: str):
    rules_data = data.splitlines()
    graph = defaultdict(list)
    for rule_data in rules_data:
        from_node, to_nodes = rule_data.split(": ")
        graph[from_node].extend(to_nodes.split(' '))
    return graph

# simply dfs every path from 'you' until we reach 'out'
# this is what I initially used for part 1
def count_paths_dfs(data: str, start = "you", end = "out"):
    graph = build_graph(data)

    paths_to_end = 0
    stack = [start]     # currently active paths
    while stack:
        node = stack.pop()
        if node == end:
            # we reached end, no need to check further
            paths_to_end += 1
            continue
        # every node connected to current node is a potential path
        stack.extend(graph[node])
    
    return paths_to_end

# memoized version of counting paths from start to end
def count_paths_memo(graph: defaultdict[list], start = "you", end = "out"):
    # I used to have some code to check for cycles, but thankfully it wasn't needed

    # its pretty much same logic as count_paths_dfs, except recursive and with memoization
    memo = {}
    def backtrack(curr: str):
        if curr in memo:
            return memo[curr]        
        if curr == end:
            ans = 1
        else:
            ans = sum(backtrack(conn) for conn in graph[curr])
        memo[curr] = ans
        return ans
    
    return backtrack(start)

# Optimized part 1 using memoization (not necessary, but faster)
def part1_with_memoiz(data: str, start = "you", end = "out"):
    graph = build_graph(data)
    return count_paths_memo(graph, start, end)

# Part 2 requires counting paths from "svr" to "out", 
#   including only those paths that go through both "fft" and "dac"
# Instead of adding booleans / int to my memoization function, 
#   I chose to compute all path arrangements and add them up
def part2(data: str):
    graph = build_graph(data)

    # getting from start to one requirement
    svr_to_fft = count_paths_memo(graph, "svr", "fft")
    svr_to_dac = count_paths_memo(graph, "svr", "dac")
    # getting from one requirement to the other
    fft_to_dac = count_paths_memo(graph, "fft", "dac")
    dac_to_fft = count_paths_memo(graph, "dac", "fft")
    # the final push to out
    fft_to_out = count_paths_memo(graph, "fft", "out")
    dac_to_out = count_paths_memo(graph, "dac", "out")

    # total paths = (start -> fft -> dac -> out) + (start -> dac -> fft -> out)
    svr_to_out = (
        svr_to_dac * dac_to_fft * fft_to_out +
        svr_to_fft * fft_to_dac * dac_to_out
    )
    return svr_to_out

sample_p1 = """aaa: you hhh
you: bbb ccc
bbb: ddd eee
ccc: ddd eee fff
ddd: ggg
eee: out
fff: out
ggg: out
hhh: ccc fff iii
iii: out"""
assert (t := count_paths_dfs(sample_p1)) == 5, f"Expected 5, got {t}"
assert (t := part1_with_memoiz(sample_p1)) == 5, f"Expected 5, got {t}"

sample_p2 = """svr: aaa bbb
aaa: fft
fft: ccc
bbb: tty
tty: ccc
ccc: ddd eee
ddd: hub
hub: fff
eee: dac
dac: fff
fff: ggg hhh
ggg: out
hhh: out"""
assert (t := part2(sample_p2)) == 2, f"Expected 2, got {t}"