| 12345678910111213141516171819202122232425262728293031323334353637 |
- #!/usr/bin/env python3
- import sys
- def read_graph(f):
- graph = {}
- for line in f:
- planet, moon = line.rstrip().split(')')
- graph.setdefault(planet, []).append(moon)
- return graph
- def count_orbits(graph, root):
- def dfs(node, depth, vis):
- vis.add(node)
- return depth + sum(dfs(nb, depth + 1, vis)
- for nb in graph.get(node, [])
- if nb not in vis)
- return dfs(root, 0, set())
- def shortest_path(graph, a, b):
- parents = {}
- for planet, moons in graph.items():
- for moon in moons:
- parents[moon] = planet
- def find_root(node):
- yield node
- if node in parents:
- yield from find_root(parents[node])
- path_a = {node: dist for dist, node in enumerate(find_root(a))}
- for dist, node in enumerate(find_root(b)):
- if node in path_a:
- return path_a[node] + dist
- graph = read_graph(sys.stdin)
- print(count_orbits(graph, 'COM'))
- print(shortest_path(graph, 'YOU', 'SAN') - 2)
|
