aoc-2022/7.py

lines = [line.strip().split(" ") for line in open("7.input").readlines()]

def insert_node(root, path, name, node):
    current = root
    for part in path:
        current = current["children"][part]

    current["children"][name] = node


def parse(lines):
    root = {"children": {}}
    pwd = []

    for line in lines:
        if line[0] == "$":
            if line[1] == "cd":
                if line[2] == "/":
                    pwd = []
                elif line[2] == "..":
                    pwd = pwd[:-1]
                else:
                    pwd.append(line[2])
            elif line[1] == "ls":
                pass
        elif line[0] == "dir":
            insert_node(root, pwd, line[1], {"children": {}})
        else:
            insert_node(root, pwd, line[1], {"size": int(line[0])})

    return root


def calc_size(node):
    if "size" not in node:
        node["size"] = sum(calc_size(child) for child in node["children"].values())

    return node["size"]


def all_dirs(node, name):
    if "children" in node:
        return [(node['size'], name)] + \
               [dir for name, child in node["children"].items() for dir in all_dirs(child, name)]
    else:
        return []


def part1(root):
    return sum(size for size, name in all_dirs(root, "/") if size <= 100000)


def part2(root):
    total = 70000000
    required = 30000000
    used = root["size"]
    to_delete = required - (total - used)

    dirs = all_dirs(root, "/")

    for size, _ in sorted(dirs):
        if size >= to_delete:
            return size

    return 0

root_node = parse(lines)
calc_size(root_node)
print(part1(root_node))
print(part2(root_node))