advent-of-code-2019/day06.py

#!/usr/bin/env python3

from lib import Graph

def preproc(puzzle_input):
    graph = Graph()
    for line in puzzle_input.split('\n'):
        sun, planet = line.split(')')
        graph.add_edge(sun, planet)
    return graph

def partI(graph):
    return sum(w for node, w in graph.min_paths("COM").items())

def partII(graph):
    return graph.min_path("SAN", "YOU") - 2

from main import Tests
tests = Tests()
tests.add("""COM)B
B)C
C)D
D)E
E)F
B)G
G)H
D)I
E)J
J)K
K)L""", partI=42)

tests.add("""COM)B
B)C
C)D
D)E
E)F
B)G
G)H
D)I
E)J
J)K
K)L
K)YOU
I)SAN""", partII=4)