day7.py

from collections import deque
from dataclasses import dataclass
import re
from typing import List, Optional, Tuple

FILE_METADATA_PATTERN = re.compile('^(?P<filesize>\d+)\s+(?P<filename>.+)')
DIR_METADATA_PATTERN = re.compile('^dir\s+(?P<dirname>.+)')


@dataclass
class File:
    dir: 'Dir'
    name: str
    size: int


@dataclass
class Dir:
    name: str
    parent_dir: Optional['Dir']
    subdirs: List['Dir']
    files: List[File]

    def add_file(self, file: File) -> None:
        if not self.has_file(file.name):
            self.files.append(file)

    def add_subdir(self, dir: 'Dir') -> None:
        if not self.has_subdir(dir.name):
            self.subdirs.append(dir)

    def get_subdir(self, dirname: str) -> Optional['Dir']:
        for dir in self.subdirs:
            if dir.name == dirname:
                return dir

        return None

    @property
    def size(self) -> int:
        ret = 0
        q = deque([self])
        while q:
            dir = q.popleft()
            ret += sum([file.size for file in dir.files])
            q += dir.subdirs

        return ret

    def has_file(self, name: str) -> bool:
        return any([file.name == name for file in self.files])

    def has_subdir(self, name: str) -> bool:
        return any([dir.name == name for dir in self.subdirs])


def get_lines() -> List[str]:
    with open('input.txt', 'r') as f:
        return [line.strip() for line in f.readlines()]

def is_cd(line: str) -> bool:
    return line.startswith('$ cd')

def parse_cd(line: str) -> str:
    assert is_cd(line)
    return line[len('$ cd '):]

def is_ls(line: str) -> bool:
    return line.startswith('$ ls')

def is_file_metadata(line: str) -> bool:
    return FILE_METADATA_PATTERN.match(line) is not None

def parse_file_metadata(line: str) -> Tuple[str, int]:
    assert is_file_metadata(line)
    m = FILE_METADATA_PATTERN.match(line)
    return (m.group('filename'), int(m.group('filesize')))

def is_dir_metadata(line: str) -> bool:
    return DIR_METADATA_PATTERN.match(line) is not None

def parse_dir_metadata(line: str) -> str:
    assert is_dir_metadata(line)
    return DIR_METADATA_PATTERN.match(line).group('dirname')

def parse() -> Dir:
    """Parse input and give root dir."""

    root = None

    curr_dir = None
    lines = get_lines()
    for line in lines:
        if is_cd(line):
            next_dir_name = parse_cd(line)
            if next_dir_name == '..':
                curr_dir = curr_dir.parent_dir if curr_dir.parent_dir else curr_dir
                continue

            if curr_dir:
                if not curr_dir.has_subdir(next_dir_name):
                    next_dir = Dir(name=next_dir_name, parent_dir=curr_dir, subdirs=[], files=[])
                    curr_dir.add_subdir(next_dir)
                else:
                    next_dir = curr_dir.get_subdir(next_dir_name)
            else:
                next_dir = Dir(name=next_dir_name, parent_dir=curr_dir, subdirs=[], files=[])

            if next_dir_name == '/':
                root = next_dir
            curr_dir = next_dir

            continue

        if is_ls(line):  # can actually just ignore
            continue

        if is_file_metadata(line):
            filename, filesize = parse_file_metadata(line)
            file = File(dir=curr_dir, name=filename, size=filesize)
            curr_dir.add_file(file)
            continue

        if is_dir_metadata(line):
            dirname = parse_dir_metadata(line)
            if curr_dir and not curr_dir.has_subdir(dirname):
                subdir = Dir(name=dirname, parent_dir=curr_dir, subdirs=[], files=[])
                curr_dir.add_subdir(subdir)

            continue

    return root

def part1():
    root = parse()
    ret = 0
    q = deque([root])
    while q:
        dir = q.popleft()
        dir_size = dir.size
        if dir_size <= 100000:
            ret += dir_size

        q += dir.subdirs

    return ret


def part2():
    total_disk_size = 70000000
    needed_disk_size = 30000000

    root = parse()
    used_disk_size = root.size
    unused_disk_size = total_disk_size - used_disk_size
    smallest = None
    q = deque([root])
    while q:
        dir = q.popleft()
        dir_size = dir.size
        if unused_disk_size + dir_size >= needed_disk_size:
            if smallest is None or dir_size < smallest:
                smallest = dir_size

        q += dir.subdirs

    return smallest