day6.py

from collections import defaultdict
from typing import NamedTuple, List

import aoc


class Orbit(NamedTuple):
    origin: str
    name: str


class OrbitTree:
    def __init__(self, name, origin=None):
        self.name = name
        self.origin = origin
        self.orbits: List["OrbitTree"] = []

    def add_orbit(self, name):
        o = OrbitTree(name, self)
        self.orbits.append(o)
        return o

    def find(self, name):
        if self.name == name:
            return self

        for o in self:
            found = o.find(name)
            if found is not None:
                return found
        return None

    def count_orbits(self, name):
        if self.name == name:
            return 0

        for o in self:
            found = o.count_orbits(name)
            if found >= 0:
                return found + 1
        return -1

    def __repr__(self):
        return f"OrbitTree(name={self.name!r}, origin={(self.origin.name if self.origin else None)!r}, orbits={[o.name for o in self.orbits]!r})"

    def __len__(self):
        return len(self.orbits)

    def __iter__(self):
        return iter(self.orbits)


def parse_orbits(data):
    def add_orbits(tree: OrbitTree, orbit: List[Orbit]):
        for obt in orbit:
            orbit_tree = tree.add_orbit(obt.name)
            add_orbits(orbit_tree, orbits[obt.name])

    orbits = defaultdict(list)
    for d in data.splitlines():
        orbit = Orbit(*d.split(')'))
        orbits[orbit.origin].append(orbit)

    o = orbits['COM']
    com = OrbitTree('COM')
    add_orbits(com, o)

    return com


def print_orbit(ot):
    print(ot)
    for o in ot:
        print_orbit(o)


def main():
    with aoc.open_input(__name__) as f:
        orbit_tree = parse_orbits(f.read())

    print(calculate_orbit_count(orbit_tree))
    print(find_santa(orbit_tree.find("YOU")))


def find_santa(you: OrbitTree, count=0):
    if you.origin is None:
        return you.count_orbits('SAN') + count

    for to_look in you.origin:
        if to_look is not you:
            santa = to_look.count_orbits('SAN')
            if santa >= 0:
                return santa + count

    return find_santa(you.origin, count + 1)


test_data = """\
COM)B
B)C
C)D
D)E
E)F
B)G
G)H
D)I
E)J
J)K
K)L
K)YOU
I)SAN"""


def test():
    orbit_tree = parse_orbits(test_data)

    print(calculate_orbit_count(orbit_tree))
    print(find_santa(orbit_tree.find("YOU")))


def calculate_orbit_count(tree: OrbitTree, count=0):
    c = count
    for t in tree:
        c += calculate_orbit_count(t, count + 1)
    return c