Advent-of-Code/2024/day9/main.py

with open('test.text', 'r') as file:
    data: str = file.read().strip()

# If block has no ID then it is free
class FreeBlock:
    size: int

    def __init__(self, s: int):
        self.size = s

    def __str__(self) -> str:
        return '.' * self.size

class FileBlock:
    size: int
    id: int

    def __init__(self, size: int, id: int):
        self.size = size
        self.id = id

    def __str__(self) -> str:
        return str(self.id) * self.size

class DiskMap:
    blocks: list[FreeBlock | FileBlock]

    def __init__(self, s: str):
        self.blocks = []
        self.parse(s)

    def parse(self, s: str) -> None:
        is_file = True
        id = 0
        for char in s:
            b = FreeBlock(int(char))
            if is_file:
                b = FileBlock(int(char), id)
                id += 1
            is_file = not is_file

            self.blocks.append(b)

    # Returns index of first free block or None
    def findFreeBlock(self) -> int | None:
        for i, block in enumerate(self.blocks):
            if type(block) == FreeBlock:
                return i

        return None

    # Returns index of last file block
    def findLastFile(self) -> int | None:
        for i, block in reversed(list(enumerate(self.blocks))):
            if type(block) == FileBlock:
                return i

        return None

    def findOffsetLastFile(self, offset: int) -> int | None:
        for i, block in list(reversed(list(enumerate(self.blocks))))[offset:]:
            if type(block) == FileBlock:
                return i

        return None

    def findFreeBlockOfBestFit(self, size: int) -> int | None:
        for i, block in enumerate(self.blocks):
            if type(block) == FreeBlock and block.size >= size:
                return i

        return None

    def moveFile(self, src: int, dest: int) -> bool:

        print('moving file')
        # Valid Input for file locations
        file: FileBlock = self.blocks[src]

        if type(file) != FileBlock:
            return False

        free: FreeBlock = self.blocks[dest]

        if type(free) != FreeBlock:
            return False

        if free.size > file.size:
            return False

        remaining_file_size = file.size - free.size
        remaining_free_size = 0

        # This means there is space left in the destination
        if remaining_file_size < 0:
            remaining_free_size = abs(remaining_file_size)
            remaining_file_size = 0

        offset = 0

        # Insert A FreeBlock after new file location if there is any free space remaining
        if remaining_free_size != 0:
            self.blocks[dest] = FileBlock(free.size - remaining_free_size, file.id)
            self.blocks.insert(dest + 1, FreeBlock(remaining_free_size))
            offset = 1

        else:
            self.blocks[dest] = FileBlock(free.size, file.id)

        self.blocks[src + offset] = FreeBlock(file.size)

        print(self)

        return True

    def compact(self) -> None:
        while True:
            free_index = self.findFreeBlock()
            file_index = self.findLastFile()

            if free_index > file_index:
                break

            free_size: int = self.blocks[free_index].size
            file_size: int = self.blocks[file_index].size
            file_id: int = self.blocks[file_index].id

            remaining_file_size = file_size - free_size
            remaining_free_size = 0
            if remaining_file_size < 0:
                remaining_free_size = abs(remaining_file_size)
                remaining_file_size = 0

            if remaining_free_size != 0:
                self.blocks[free_index] = FileBlock(free_size - remaining_free_size, file_id)
                self.blocks.insert(free_index + 1, FreeBlock(remaining_free_size))
                self.blocks[file_index + 1] = FreeBlock(file_size)
            else:
                self.blocks[free_index] = FileBlock(free_size, file_id)
                self.blocks[file_index] = FreeBlock(free_size)

            if remaining_file_size:
                self.blocks.insert(file_index, FileBlock(remaining_file_size, file_id))

    def compact2(self) -> None:
        offset = 0
        while True:
            file_index: int = self.findOffsetLastFile(offset)
            if file_index == None:
                break

            file_size: int = self.blocks[file_index].size

            free_index: int = self.findFreeBlockOfBestFit(file_size)
            if free_index == None:
                break

            free_size: int = self.blocks[free_index].size


            self.moveFile(file_index, free_index)
offset = len(self.blocks) - file_index



    def __str__(self) -> str:
        s = ''
        for block in self.blocks:
            s += str(block)
        return s

# Parse Data
diskmap = DiskMap(data)

# Part 1
print(f'Input: {diskmap}')
diskmap.compact()
print(f'Compacted: {diskmap}')

total = 0
index = 0
for block in diskmap.blocks:
    if type(block) == FileBlock:
        for i in range(block.size):
            total += block.id * index
            index += 1

print(f'Part 1: {total}')

# Part 2
diskmap = DiskMap(data)
print('Part 2')
print(f'Input: {diskmap}')
diskmap.compact2()
print(f'Compacted: {diskmap}')

total = 0
index = 0
for block in diskmap.blocks:
    if type(block) == FileBlock:
        for i in range(block.size):
            total += block.id * index
            index += 1

print(f'Part 2: {total}')