add day 6 part 2

.travis.yml

@@ -9,7 +9,7 @@ jobs:
         - cd python
       script:
         - pytest
-        - mypy 1 2 3 4 5
+        - mypy 1 2 3 4 5 6
 
     - language: rust
       rust: stable

python/6/day6.py

@@ -1,4 +1,5 @@
-from typing import Iterator, Dict, Set
+from typing import Iterator, Dict, Set, List
+
 
 Node = str
 ChildGraph = Dict[Node, Set[Node]]
@@ -17,7 +18,7 @@ def parse_input(input: str) -> ChildGraph:
     return data
 
 
-def reverse_graph(graph: ChildGraph):
+def reverse_graph(graph: ChildGraph) -> ParentGraph:
     output: ParentGraph = {}
     for parent, children in graph.items():
         for child in children:
@@ -44,7 +45,7 @@ def count_all_orbits(graph: ParentGraph) -> int:
     return orbits
 
 
-def part1():
+def part1() -> int:
     with open("6/input.txt") as f:
         input = f.read()
     gr = parse_input(input)
@@ -52,5 +53,34 @@ def part1():
     return count_all_orbits(parent_graph)
 
 
+def get_common_ancestor(a: List[Node], b: List[Node]) -> Node:
+    node = "none"
+    i = -1
+    while a[i] == b[i]:
+        node = a[i]
+        i -= 1
+    return node
+
+
+def calculate_distance_to_santa(pg: ParentGraph) -> int:
+    your_parents = list(get_indirect_orbits(pg, "YOU"))
+    santa_parents = list(get_indirect_orbits(pg, "SAN"))
+    common_ancestor = get_common_ancestor(your_parents, santa_parents)
+
+    up = your_parents.index(common_ancestor)
+    down = santa_parents.index(common_ancestor)
+
+    return up + down
+
+
+def part2() -> int:
+    with open("6/input.txt") as f:
+        input = f.read()
+    gr = parse_input(input)
+    parent_graph = reverse_graph(gr)
+    return calculate_distance_to_santa(parent_graph)
+
+
 if __name__ == '__main__':
-    print(part1())
+    # print(part1())
+    print(part2())

python/6/test_day6.py

@@ -1,4 +1,5 @@
-from day6 import parse_input, reverse_graph, count_all_orbits, get_indirect_orbits
+from day6 import parse_input, reverse_graph, count_all_orbits, get_indirect_orbits, get_common_ancestor, \
+    calculate_distance_to_santa
 
 example_input = """COM)B
 B)C
@@ -12,6 +13,8 @@ E)J
 J)K
 K)L"""
 
+extended_input = example_input + "\nK)YOU\nI)SAN"
+
 
 def test_parse_input():
     correct = {
@@ -51,3 +54,30 @@ def test_complete():
     assert list(get_indirect_orbits(parent_graph, "COM")) == []
 
     assert count_all_orbits(parent_graph) == 42
+
+
+def test_b_complete():
+    child_graph = parse_input(extended_input)
+    parent_graph = reverse_graph(child_graph)
+    your_parents = list(get_indirect_orbits(parent_graph, "YOU"))
+    assert your_parents == ['K', 'J', 'E', "D", 'C', 'B', "COM"]
+    santa_parents = list(get_indirect_orbits(parent_graph, "SAN"))
+    assert santa_parents == ["I", "D", 'C', 'B', "COM"]
+
+    common_ancestor = get_common_ancestor(your_parents, santa_parents)
+    assert common_ancestor == "D"
+
+    assert your_parents.index(common_ancestor) == 3
+    assert santa_parents.index(common_ancestor) == 1
+
+
+def test_get_common_ancestor():
+    a = ["f", "e", "d", "c", "b", "a"]
+    b = ["g", "c", "b", "a"]
+    assert get_common_ancestor(a, b) == "c"
+
+
+def test_calculate_distance_to_santa():
+    child_graph = parse_input(extended_input)
+    parent_graph = reverse_graph(child_graph)
+    assert calculate_distance_to_santa(parent_graph) == 4