2019: d06: ex2: add solution

2019/d06/ex2/ex2.py

@@ -0,0 +1,86 @@
+#!/usr/bin/env python
+
+import sys
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional
+
+
+@dataclass
+class OrbitGraph:
+    name: str
+    children: List["OrbitGraph"] = field(default_factory=list)
+
+
+def make_orbits(description: List[str]) -> OrbitGraph:
+    orbits: Dict[str, OrbitGraph] = {}
+
+    def find_or_add(name: str) -> OrbitGraph:
+        if name in orbits:
+            return orbits[name]
+        orbit = OrbitGraph(name)
+        orbits[name] = orbit
+        return orbit
+
+    for l in description:
+        parent, child = map(find_or_add, map(lambda x: x.strip(), l.split(")")))
+        parent.children.append(child)
+
+    return orbits["COM"]  # Assume common root is named 'COM'
+
+
+def count_orbits_hop(root: OrbitGraph) -> int:
+    nodes: List[OrbitGraph] = []
+
+    def ind_node(node: OrbitGraph) -> int:
+        return nodes.index(node)
+
+    def fill_nodes(root: OrbitGraph) -> None:
+        nodes.append(root)
+        for child in root.children:
+            fill_nodes(child)
+
+    fill_nodes(root)
+    n = len(nodes)
+    dist: List[List[int]] = [[2 * n] * len(nodes) for __ in range(len(nodes))]
+    next_nodes: List[List[Optional[OrbitGraph]]] = [[None] * n for __ in range(n)]
+
+    def fill_mat(root: OrbitGraph) -> None:
+        # From root to itself
+        dist[ind_node(root)][ind_node(root)] = 0
+        next_nodes[ind_node(root)][ind_node(root)] = root
+        for child in root.children:
+            # From root to child
+            dist[ind_node(root)][ind_node(child)] = 1
+            next_nodes[ind_node(root)][ind_node(child)] = child
+            # The other way
+            dist[ind_node(child)][ind_node(root)] = 1
+            next_nodes[ind_node(child)][ind_node(root)] = root
+            # Do it again
+            fill_mat(child)
+
+    fill_mat(root)
+
+    for k in range(n):
+        print(f"{k} / {n} rounds")
+        for i in range(n):
+            for j in range(n):
+                new_dist = dist[i][k] + dist[k][j]
+                if dist[i][j] > new_dist:
+                    dist[i][j] = new_dist
+                    next_nodes[i][j] = next_nodes[i][k]
+
+    useful_indices: Dict[str, int] = {}
+    for index, node in enumerate(nodes):
+        if node.name in ("YOU", "SAN"):
+            useful_indices[node.name] = index
+    return (
+        dist[useful_indices["YOU"]][useful_indices["SAN"]] - 2
+    )  # Because we are in orbit
+
+
+def main() -> None:
+    print(count_orbits_hop(make_orbits(sys.stdin.readlines())))
+
+
+if __name__ == "__main__":
+    main()