advent-of-code/2017-python/solutions/day_07.py

from solutions import BaseSolution


class Program:
    def __init__(self, data=None):
        if data:
            name, weight, disc = self._parse(data)
            self.name = name
            self.weight = weight
            self.disc = disc

    def __repr__(self):
        return str(self.name)

    def _parse(self, data):
        disc = []
        try:
            name, weight = data.split()
        except ValueError:
            name, weight, _, *disc = data.split()
        weight = int(weight[1 : len(weight) - 1])
        disc = tuple(p.replace(",", "") for p in disc)
        return name, weight, disc

    def has_disc(self):
        return len(self.disc) > 0

    def unseen_discs(self, seen_discs):
        return [t for t in self.disc if t not in seen_discs]


class Solution(BaseSolution):
    input_file = "07.txt"

    def __str__(self):
        return "Day 7: Recursive Circus"

    def _get_programs(self, puzzle_input):
        return [Program(data) for data in puzzle_input.splitlines()]

    def _get_discs(self, disc, programs):
        subdisc = [
            {"own_weight": p.weight, "obj": p} for p in programs if p.name in disc
        ]
        for t in subdisc:
            t["weight"] = t["own_weight"]
            if t["obj"].has_disc():
                t["disc"] = self._get_discs(t["obj"].disc, programs)
                t["weight"] += sum([st["weight"] for st in t["disc"]])
            del t["obj"]
        return subdisc

    def _find_unbalanced_disc(self, disc):
        disc = sorted(disc, key=lambda t: t["weight"])
        if not disc[0]["weight"] < disc[1]["weight"]:
            disc = sorted(disc, key=lambda t: t["weight"], reverse=True)
        return disc[0], disc[1]["weight"] - disc[0]["weight"]

    def solve(self, puzzle_input):
        programs = self._get_programs(puzzle_input)
        programs_with_discs = []
        seen = []
        for p in programs:
            if p.has_disc():
                programs_with_discs.append(p)
            else:
                seen.append(p.name)
        for p in programs_with_discs:
            unseen = p.unseen_discs(seen)
            if len(unseen) > 0:
                seen += unseen
        bottom_program = list(
            filter(lambda p: p.name not in seen, programs_with_discs)
        )[0]
        return bottom_program

    def solve_again(self, puzzle_input):
        programs = self._get_programs(puzzle_input)
        bottom_program = self.solve(puzzle_input)
        disc_tree = {
            "own_weight": bottom_program.weight,
            "disc": self._get_discs(bottom_program.disc, programs),
        }
        diff = -1
        unbalanced = True
        while unbalanced:
            disc, new_diff = self._find_unbalanced_disc(disc_tree["disc"])
            if new_diff == 0:
                unbalanced = False
            else:
                disc_tree = disc
                diff = new_diff
        return disc_tree["own_weight"] + diff


if __name__ == "__main__":
    solution = Solution()
    solution.show_results()