day09/solution_gemini-1.5-pro.py

file = "input.txt"

def solve_part1(disk_map):
    disk_blocks = []
    file_id = 0
    is_file = True
    for i in range(0, len(disk_map), 2):
        file_size = int(disk_map[i])
        free_size = int(disk_map[i+1])

        disk_blocks.extend([file_id] * file_size)
        disk_blocks.extend(["."] * free_size)
        file_id += 1
    
    disk_str = "".join(map(str,disk_blocks))
    
    while "." in disk_str:
        first_dot = disk_str.find(".")
        last_file_block = disk_str.rfind(str(file_id - 1))
        
        disk_list = list(disk_str)
        disk_list[first_dot] = disk_list[last_file_block]
        disk_list[last_file_block] = "."
        disk_str = "".join(disk_list)

    checksum = 0
    for i in range(len(disk_str)):
        if disk_str[i] != ".":
            checksum += i * int(disk_str[i])
    return checksum

def solve_part2(disk_map):
    disk_blocks = []
    file_sizes = {}
    file_id = 0
    is_file = True
    for i in range(0, len(disk_map), 2):
        file_size = int(disk_map[i])
        free_size = int(disk_map[i+1])
        file_sizes[file_id] = file_size
        disk_blocks.extend([file_id] * file_size)
        disk_blocks.extend(["."] * free_size)
        file_id += 1

    disk_str = "".join(map(str, disk_blocks))

    for fid in range(file_id - 1, -1, -1):
        file_size = file_sizes[fid]
        first_occurrence = disk_str.find(str(fid))
        free_space_size = 0
        free_space_start = -1

        for i in range(len(disk_str)):
            if disk_str[i] == ".":
                if free_space_start == -1:
                    free_space_start = i
                free_space_size += 1
            else:
                if free_space_size >= file_size and free_space_start < first_occurrence:
                    disk_list = list(disk_str)
                    for j in range(file_size):
                        disk_list[free_space_start + j] = str(fid)
                        disk_list[first_occurrence + j] = "."
                    disk_str = "".join(disk_list)
                    break
                free_space_start = -1
                free_space_size = 0
        
    checksum = 0
    for i in range(len(disk_str)):
        if disk_str[i] != ".":
            checksum += i * int(disk_str[i])
    return checksum

with open(file) as f:
    disk_map = f.readline().strip()

result1 = solve_part1(disk_map)
print(result1)

result2 = solve_part2(disk_map)
print(result2)