advent-of-code/2019/d06/ex1/ex1.py

54 lines
1.3 KiB
Python
Raw Normal View History

2019-12-08 22:13:54 +01:00
#!/usr/bin/env python
import sys
from dataclasses import dataclass, field
from typing import Dict, List
@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(root: OrbitGraph) -> int:
ans = 0
stack = 0 # Count the number of direct and indirect orbits to the current root
def helper(root: OrbitGraph):
nonlocal ans
nonlocal stack
ans += stack # Count the number of orbits to this node
stack += 1 # Add the current root to stack of parents
for child in root.children:
helper(child) # Count those orbits for each child
stack -= 1 # Remove the current root from the stack of parents
helper(root)
return ans
def main() -> None:
print(count_orbits(make_orbits(sys.stdin.readlines())))
if __name__ == "__main__":
main()