day_7b.py

from __future__ import annotations

from aocd import lines, numbers


class Node:
    def __init__(self, name, type, size=0):
        self.name = name
        self.type = type
        self.size = size
        self.children = []

    def __repr__(self):
        return f"{self.name=}, {self.type=}, {self.size=}, {self.children=}"

    def total_size(self):
        val = self.size
        for n in self.children:
            val += n.total_size()
        return val


stack = []
nodes = []
for line in lines:
    if line == "$ cd ..":
        stack.pop()
        continue
    elif line.startswith("$ cd"):
        name = line[5:]
        node = Node(name, type="dir")
        if stack:
            stack[-1].children.append(node)
        stack.append(node)
        nodes.append(node)
        continue
    elif line.startswith("$ ls"):
        continue
    elif line.startswith("dir"):
        continue
    else:
        size, name = line.split()
        node = Node(name=name, type="file", size=int(size))
        stack[-1].children.append(node)
        nodes.append(node)

curr = nodes[0].total_size()
free = 70000000 - curr
needed = 30000000 - free
print(sorted((n.total_size() for n in nodes if n.total_size() >= needed and n.type == "dir")))