Implement 2019 day 6

2019/aoc2019/day06.py

@@ -0,0 +1,27 @@
+from typing import TextIO
+
+import networkx
+
+
+def read_graph(data: TextIO) -> networkx.DiGraph:
+    graph = networkx.DiGraph()
+
+    for line in data:
+        a, b = line.strip().split(')')
+        graph.add_edge(a, b)
+
+    return graph
+
+
+def part1(data: TextIO) -> int:
+    graph = read_graph(data)
+
+    paths = networkx.single_source_shortest_path_length(graph, 'COM')
+
+    return sum(paths.values())
+
+
+def part2(data: TextIO) -> int:
+    graph = read_graph(data).to_undirected()
+
+    return networkx.shortest_path_length(graph, 'YOU', 'SAN') - 2

2019/requirements.txt

@@ -1 +1,2 @@
 pytest
+networkx

2019/tests/test_day6.py

@@ -0,0 +1,39 @@
+from io import StringIO
+
+from aoc2019.day06 import part1, part2
+
+
+def test_sample_part1():
+    data = StringIO("""\
+            COM)B
+            B)C
+            C)D
+            D)E
+            E)F
+            B)G
+            G)H
+            D)I
+            E)J
+            J)K
+            K)L""")
+
+    assert part1(data) == 42
+
+
+def test_sample_part2():
+    data = StringIO("""\
+            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""")
+
+    assert part2(data) == 4