Advent of Code 2024, day 09

AdventOfCode2024.Day09.py

with open("./data/input09.txt") as inputFile:
    input = inputFile.read()

fs = []
total_free_blocks = 0
i = 0
while i < len(input):
    file_id = i // 2
    file_size = int(input[i])
    free_blocks = 0 if i == len(input)-1 else int(input[i+1])
    fs.append((file_id, file_size, free_blocks))
    total_free_blocks += free_blocks
    i += 2

# Part 1

fs_dup = fs.copy()
defragmented = []
cur_pos = 0
checksum = 0

def move_blocks(file_id, block_size, checksum, cur_pos):
    defragmented.append((file_id, block_size))
    checksum += ((2 * cur_pos + block_size - 1) * block_size) // 2 * file_id
    cur_pos += block_size
    return (checksum, cur_pos)

ndxl = 0
ndxr = len(fs_dup) - 1
while ndxr >= ndxl:
    file_id_l, file_size_l, free_blocks_l = fs_dup[ndxl]
    if file_size_l > 0:
        checksum, cur_pos = move_blocks(file_id_l, file_size_l, checksum, cur_pos)
        fs_dup[ndxl] = (file_id_l, 0, free_blocks_l)
    else:
        file_id_r, file_size_r, free_blocks_r = fs_dup[ndxr]
        if free_blocks_l >= file_size_r:
            checksum, cur_pos = move_blocks(file_id_r, file_size_r, checksum, cur_pos)
            fs_dup[ndxl] = (file_id_l, file_size_l, free_blocks_l - file_size_r)
            ndxr -= 1
            if free_blocks_l == file_size_r: ndxl += 1
        else:
            checksum, cur_pos = move_blocks(file_id_r, free_blocks_l, checksum, cur_pos)
            fs_dup[ndxr] = (file_id_r, file_size_r - free_blocks_l, free_blocks_r)
            ndxl += 1

print(checksum)

# Part 2

regions = []
for elt in fs:
    file_id, file_size, free_size = elt
    regions.append((file_id, file_size))
    regions.append((None, free_size))

for ndx in range(len(regions)):
    ndxr = len(regions) - 1 - ndx
    file_id_r, file_size_r = regions[ndxr]
    if file_id_r != None:
        for ndxl in range(ndxr):
            file_id_l, block_size_l = regions[ndxl]
            if file_id_l == None and block_size_l >= file_size_r:
                regions[ndxl] = (file_id_r, file_size_r)
                if block_size_l > file_size_r:
                    regions.insert(ndxl+1, (None, block_size_l - file_size_r))
                    regions[ndxr+1] = (None, file_size_r)
                else:
                    regions[ndxr] = (None, file_size_r)
                break

cur_pos = 0
checksum = 0
for ndx in range(len(regions)):
    file_id, file_size = regions[ndx]
    if file_id != None:
        checksum += ((2 * cur_pos + file_size - 1) * file_size) // 2 * file_id
    cur_pos += file_size
print(checksum)