aoc/2022/7/solution.py

#!/usr/bin/python3
import sys
import re

# part 1

def print_tree(tree, start='/', level=0):
    subdirs, size = tree
    indent='\t' * level
    print(f'{indent}{start} ({size}):')
    for d in subdirs:
        print_tree(tree[0][d], f'{d}', level + 1)

def update_dirs(tree, pwd):
    current_branch = tree[0]
    for d in pwd.lstrip('/').split('/'):
        if d not in current_branch:
            current_branch[d] = [{}, 0]
        current_branch = current_branch[d][0]

def update_sizes(tree, pwd, size):
    #print('updating size for pwd', pwd, size)
    current_branch = tree
    tree[1] += size
    for d in pwd.lstrip('/').split('/'):
        if d in current_branch[0]:
            current_branch = current_branch[0][d]
            current_branch[1] += size

def recurse_tree_for_sizes(tree):
    total = 0
    subdirs, size = tree
    total += size if size <= 100000 else 0
    for d in subdirs:
        total += recurse_tree_for_sizes(tree[0][d])
    return total

def recurse_tree_for_deletion(tree, candidates, rootsize):
    subdirs, size = tree
    if 70000000 - rootsize + size >= 30000000:
        candidates.append(size)
    for d in subdirs:
        recurse_tree_for_deletion(tree[0][d], candidates, rootsize)

def solution():
    tree = [{}, 0]
    pwd = None
    for line in sys.stdin:
        l = line.strip()
        match = None
        if (match := re.match(r'^\$ cd (.*)', l)) is not None:
            target = match.groups()[0]
            #print(pwd, target)
            if target == '/':
                pwd = target
            elif target == '..':
                index = pwd.rfind('/')
                pwd = '/' if index == 0 else pwd[:index]
            else:
                pwd += ('/' if len(pwd) > 1 else '') + target
                update_dirs(tree, pwd)
            continue
        elif re.match(r'^\$ ls', l):
            continue
        elif (match := re.match(r'^dir (.*)', l)) is not None:
            continue
        elif (match := re.match(r'^([0-9]+) (.*)', l)) is not None:
            size, file = match.groups()
            #print('size', size, 'file', file)
            update_sizes(tree, pwd, int(size))
            continue
        else:
            continue
    #print(tree)
    print_tree(tree)
    print(f'total sizes come to {recurse_tree_for_sizes(tree)}')

    # part 2
    candidates = []
    recurse_tree_for_deletion(tree, candidates, tree[1])
    print(f'total sizes come to {min(candidates)}')

solution()