87 lines
2.5 KiB
Python
87 lines
2.5 KiB
Python
|
#!/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()
|