AdventOfCode/2022/07/a.py

import os, argparse, re


def solve_task(lines):
    e = Explorer()
    for line in lines:
        e.parse_output(line)

    folders = e.get_all_folders()
    small_folders = [f for f in folders if f.size < 100000]
    small_folders_size = sum([f.size for f in small_folders])
    print(f"Total size: {e.get_size()}, sum of all small folder sizes: {small_folders_size}")

class Explorer:
    def __init__(self):
        self.root_node = FolderNode()
        self.current_dir = self.root_node
    
    def parse_output(self, line):
        tokens = line.split()
        if tokens[0] == "$":
            self.parse_cmd(tokens[1:])
        elif tokens[0] == "dir":
            self.current_dir.add_folder(tokens[1])
        else:
            self.current_dir.add_file(tokens[1], tokens[0])

    def get_size(self):
        return self.root_node.size
    
    def parse_cmd(self, tokens):
        cmd = tokens[0]
        if cmd == "cd":
            path = tokens[1]
            if path == "..":
                self.current_dir = self.current_dir.parent
            elif path == "/":
                self.current_dir = self.root_node
            else:
                if path in self.current_dir.children:
                    # Existing name
                    self.current_dir = self.current_dir.children[path]
                else:
                    self.current_dir.add_folder(path)
                    self.current_dir = self.current_dir.children[path]
    
    def get_all_folders(self):
        all_folders_found = [self.root_node]
        folders_to_traverse = [self.root_node]
        while len(folders_to_traverse) > 0:
            current_folder = folders_to_traverse.pop()
            folders_found = [folder for folder in current_folder.children.values() if isinstance(folder, FolderNode)]
            all_folders_found.extend(folders_found)
            folders_to_traverse.extend(folders_found)
        return all_folders_found

class Node:
    def __init__(self, parent):
        self.size = 0
        if parent:
            self.parent = parent
        else:
            self.parent = self

class FileNode(Node):
    def __init__(self, parent, name, size):
        super().__init__(parent)
        self.parent = parent
        self.name = name
        self.size = int(size)

class FolderNode(Node):
    def __init__(self, parent = None, name = "/"):
        super().__init__(parent)
        self.children = {}
        self.name = name

    def add(self, node):
        self.children[node.name] = node
        self.adjust_size(node.size)
    
    def add_file(self, name, size):
        file = FileNode(self, name, size)
        self.add(file)

    def add_folder(self, name):
        folder = FolderNode(self, name)
        self.add(folder)
    
    def adjust_size(self, increment):
        self.size += increment
        if not self.is_root():
            self.parent.adjust_size(increment)

    def is_root(self):
        return self.parent == self


def read_lines(filename):
    lines = []
    with open(filename) as infile:
        for raw_line in infile:
            line = raw_line.rstrip()
            lines.append(line)
    return lines

def parse_arguments():
    parser = argparse.ArgumentParser(description="Script that solves the case",epilog="Have a nice day!")
    parser.add_argument('filename', nargs='?', default="example.txt", help='Input file')
    args = parser.parse_args()
    return args

def main():
    args = parse_arguments()
    lines = read_lines(args.filename)
    solve_task(lines)

if __name__ == "__main__":
    main()