Spaces:
Running
Running
| def load_data(file): | |
| with open(file) as f: | |
| return f.read().strip("\n") | |
| def generate_blocks(disk_map): | |
| blocks = [] | |
| for idx, i in enumerate(disk_map): | |
| if idx % 2: | |
| blocks.extend( ["."]*int(i)) | |
| else: | |
| blocks.extend( [str(idx // 2)] *int(i)) | |
| return blocks | |
| def get_next_valid_block_from_end(end: int, blocks): | |
| for j in range(end, -1, -1): | |
| if blocks[j] != ".": | |
| return j | |
| def compact_blocks(blocks): | |
| # blocks = list(blocks) | |
| end = get_next_valid_block_from_end(len(blocks) - 1, blocks) | |
| for idx, block in enumerate(blocks): | |
| if end == idx: | |
| break | |
| if block == ".": | |
| blocks[idx] = blocks[end] | |
| blocks[end] = "." | |
| end = get_next_valid_block_from_end(end, blocks) | |
| # print("".join(blocks)) | |
| return blocks | |
| def compute_checksum(blocks): | |
| checksum = 0 | |
| for idx, num in enumerate(blocks): | |
| if num == ".": | |
| return checksum | |
| checksum += idx*int(num) | |
| return checksum | |
| disk_map = load_data("input.txt") | |
| blocks = generate_blocks(disk_map) | |
| compact_blocks = compact_blocks(blocks) | |
| print(compute_checksum(compact_blocks)) | |
| ## Part two | |
| def generate_blocks(disk_map): | |
| blocks = [] | |
| for idx, i in enumerate(disk_map): | |
| if idx % 2: | |
| if int(i) > 0: | |
| blocks.append( ["."]*int(i)) | |
| else: | |
| blocks.append( [str(idx // 2)] *int(i)) | |
| return blocks | |
| def get_free_space_map(blocks): | |
| free_space_map = [] | |
| pos = 0 | |
| for idx, block in enumerate(blocks): | |
| if block[0] == ".": | |
| free_space_map.append((idx, len(block))) | |
| pos += len(block) | |
| return free_space_map | |
| def get_next_valid_block_from_end(blocks): | |
| for idx, block in enumerate(blocks[::-1]): | |
| if block[0] != ".": | |
| block_idx = len(blocks) - idx - 1 | |
| return block, block_idx | |
| def pprint(blocks): | |
| print("".join(["".join(b) for b in blocks])) | |
| def move_block_at_fileid(blocks, file_id: int): | |
| free_space_map = get_free_space_map(blocks) | |
| # block = blocks[block_idx] | |
| for idx, block in enumerate(blocks): | |
| if block[0] != "." and block[0] == str(file_id): | |
| block_idx = idx | |
| break | |
| K = len(block) | |
| for free_block_idx, free_len in free_space_map: | |
| if free_len > K and free_block_idx < block_idx: | |
| # First condition means we have more space than needed | |
| # Second condition ensures we Only move forward in queue, not backwards | |
| blocks[free_block_idx] = ["."] * (free_len - len(block)) # Remaining free memory after insert | |
| blocks.insert(free_block_idx, block) # Insert at index | |
| blocks[block_idx + 1] = ["."] * K # Overwrite at end (+1 because we added a new index) | |
| return True | |
| elif free_len == K and free_block_idx < block_idx: | |
| # First condition means we have just enough space | |
| # Second condition ensures we Only move forward in queue, not backwards | |
| blocks[free_block_idx] = block | |
| blocks[block_idx] = ["."] * K | |
| return True | |
| return False | |
| def compute_checksum(blocks): | |
| blocks_ext = [b for block in blocks for b in block] | |
| checksum = 0 | |
| for idx, num in enumerate(blocks_ext): | |
| if num == ".": | |
| continue | |
| checksum += idx*int(num) | |
| return checksum | |
| def get_max_file_id(blocks): | |
| for block in blocks[::-1]: | |
| if block[0] != ".": | |
| max_file_id = int(block[0]) | |
| return max_file_id | |
| disk_map = load_data("input.txt") | |
| blocks = generate_blocks(disk_map) | |
| max_file_id = get_max_file_id(blocks) | |
| # pprint(blocks) | |
| for file_id in range(max_file_id, -1, -1): | |
| # Progress-bar | |
| # if file_id % 1000 == 0: | |
| # print(file_id) | |
| moved = move_block_at_fileid(blocks, file_id) | |
| # if moved: | |
| # pprint(blocks) | |
| print(compute_checksum(blocks)) | |