aoc2024day9.py

# === Parsing ===

with open("input.txt", "r") as f:
    filesystem_string = f.read().strip()

filesystem = [
    [(idx // 2 if idx % 2 == 0 else None, int(char))]
    for idx, char in enumerate(filesystem_string)
]

print(filesystem)


# === Part 1 ===
from itertools import chain


def compact_filesystem(filesystem):
    left_ptr = 1
    right_ptr = len(filesystem) - 1
    # We move files while there's space to the left of the
    # rightmost file.
    while left_ptr < right_ptr:
        _, space = filesystem[left_ptr].pop()
        idx, length = filesystem[right_ptr].pop()
        # How much of the file can we move?
        to_move = min(space, length)
        # Move that portion of the file.
        filesystem[left_ptr].append((idx, to_move))

        # Is there any space left where we moved the file to?
        if to_move == space:
            left_ptr += 2
        elif to_move < space:
            filesystem[left_ptr].append((None, space - to_move))

        # Did we move the whole file to the empty space?
        if to_move == length:
            filesystem.pop()  # Get rid of the original file sublist.
            if left_ptr < right_ptr:
                filesystem.pop()  # Get rid of the empty space to its left.
            right_ptr -= 2
        elif to_move < length:
            filesystem[right_ptr].append((idx, length - to_move))

    return filesystem


def print_fs(filesystem):
    for idx, length in chain.from_iterable(filesystem):
        if idx is None:
            print("." * length, end="")
        else:
            print(f"{idx}" * length, end="")
    print()


def checksum(filesystem):
    abs_position = 0
    total = 0
    for idx, length in chain.from_iterable(filesystem):
        for _ in range(length):
            total += abs_position * idx
            abs_position += 1
    return total


print_fs(filesystem)
filesystem = compact_filesystem(filesystem)
print_fs(filesystem)
print(checksum(filesystem))


# === Part 2 ===

with open("input.txt", "r") as f:
    filesystem_string = f.read().strip()

filesystem = [
    [(idx // 2 if idx % 2 == 0 else None, int(char))]
    for idx, char in enumerate(filesystem_string)
]


def compact_contiguous_filesystem(filesystem):
    # Try to move a file.
    for right_ptr in range(len(filesystem) - 1, -1, -2):
        idx, length = filesystem[right_ptr].pop()
        # Look through each possibly empty space.
        for left_ptr in range(1, right_ptr, 2):
            # Try to move the file...
            possible_index, space = filesystem[left_ptr][-1]
            if possible_index is None and length <= space:
                # The file can be moved here.
                filesystem[left_ptr].pop()
                filesystem[left_ptr].append((idx, length))
                filesystem[right_ptr].append((None, length))
                # Did the file fill up the whole space?
                if space > length:
                    filesystem[left_ptr].append((None, space - length))
                break
        else:  # We did not find a suitable empty space
            filesystem[right_ptr].append((idx, length))

    return filesystem


def updated_checksum(filesystem):
    abs_position = 0
    total = 0
    for idx, length in chain.from_iterable(filesystem):
        for _ in range(length):
            if idx is not None:
                total += abs_position * idx
            abs_position += 1
    return total


filesystem = compact_contiguous_filesystem(filesystem)
# print_fs(filesystem)
print(updated_checksum(filesystem))