import time
start_time = time.time()

with open('puzzle.txt', 'r') as f:
    line = f.read()

disk_string = []

block_length = {}
block_initial_index = {}

file = True

file_id = -1

for char in list(line):
    if file:
        file_id += 1
        disk_string += [str(file_id)] * int(char)
        block_length[str(file_id)] = int(char)
        block_initial_index[str(file_id)] = len(disk_string) - block_length[str(file_id)]

    else:
        disk_string += ['.'] * int(char)

    file = not file

while file_id != 0:
    file_id = str(file_id)

    range_a_found = False

    for index_a, char_a in enumerate(disk_string):
        if char_a == '.':
            empty_count = 1
            while index_a + empty_count < block_initial_index[file_id] and disk_string[index_a + empty_count] == '.':
                empty_count += 1

            if empty_count >= block_length[file_id] and index_a < block_initial_index[file_id]:
                for i in range(block_length[file_id]):
                    disk_string[index_a + i] = file_id
                    disk_string[block_initial_index[file_id] + i] = '.'

                break

    file_id = int(file_id)
    file_id -= 1

checksum = 0

for index, char in enumerate(disk_string):
    if char != '.':
        checksum += int(char) * index

print(checksum)
print(f'Total time: {time.time() - start_time}')