My Advent of Code solutions in Python.
kevinyap.ca/2019/12/going-fast-in-advent-of-code/
advent-of-code
python
Add solution for 2024/09
+117
+117
2024/day09.py
+117
2024/day09.py
···
1
+
import copy
2
+
import fileinput
3
+
4
+
5
+
def read_disk(disk_map, part_2=False):
6
+
"""
7
+
Takes a disk map as a string and returns a list of the form
8
+
(block_id, length) and the maximium block ID. For free space,
9
+
`block_id` is set to None.
10
+
"""
11
+
# Pad the disk with free space at the end for part 1.
12
+
disk_map += '0'
13
+
14
+
disk = []
15
+
for block_id, (file_length, free_length) in enumerate(zip(disk_map[::2], disk_map[1::2])):
16
+
if part_2:
17
+
disk.append((block_id, int(file_length)))
18
+
disk.append((None, int(free_length)))
19
+
else:
20
+
for _ in range(int(file_length)):
21
+
disk.append((block_id, 1))
22
+
for _ in range(int(free_length)):
23
+
disk.append((None, 1))
24
+
25
+
return disk, block_id
26
+
27
+
def disk_checksum(disk):
28
+
"""Returns the checksum of a disk."""
29
+
checksum = 0
30
+
idx = 0
31
+
for block_id, length in disk:
32
+
if block_id == None:
33
+
idx += length
34
+
else:
35
+
for _ in range(length):
36
+
checksum += block_id * idx
37
+
idx += 1
38
+
39
+
return checksum
40
+
41
+
42
+
def compact_disk(disk, max_block_id, part_2=False):
43
+
"""Performs the compacting operation on the disk."""
44
+
disk = copy.deepcopy(disk)
45
+
46
+
# The current block we should try to shift (if in Part 2).
47
+
block = max_block_id
48
+
49
+
while True:
50
+
# Search for the next block to compact.
51
+
idx = len(disk) - 1
52
+
53
+
# If part 2, it's when we find a block matching the expected ID.
54
+
if part_2:
55
+
while disk[idx][0] != block:
56
+
idx -= 1
57
+
if idx < 0:
58
+
break
59
+
60
+
# Otherwise, it's whenever we find a non-empty block.
61
+
else:
62
+
while disk[idx][0] == None:
63
+
idx -= 1
64
+
if idx < 0:
65
+
break
66
+
67
+
# Try to perform a shift with the left-most free space we find.
68
+
block_id, block_length = disk[idx]
69
+
for free in range(len(disk)):
70
+
free_id, free_length = disk[free]
71
+
72
+
# If the `free` pointer is to the right of the `idx` pointer,
73
+
# we've moved too far to attempt to find a possible swap.
74
+
if idx < free:
75
+
break
76
+
77
+
# We've found a valid block to move!
78
+
if free_id == None and block_length <= free_length:
79
+
# Delete the existing free and file blocks (order matters).
80
+
del disk[idx]
81
+
del disk[free]
82
+
83
+
# Create new free and file blocks (order matters).
84
+
disk.insert(free, (block_id, block_length))
85
+
disk.insert(idx, (None, block_length))
86
+
87
+
# If we're in part 2 and we moved the file into a free
88
+
# spot with more room available, we need to create a
89
+
# new free block with the remainder space.
90
+
if block_length < free_length:
91
+
disk.insert(free + 1, (None, free_length - block_length))
92
+
93
+
break
94
+
95
+
# Termination conditions.
96
+
if part_2:
97
+
block -= 1
98
+
if block == 0:
99
+
break
100
+
else:
101
+
if idx < free:
102
+
break
103
+
104
+
return disk
105
+
106
+
107
+
# Read problem input.
108
+
DISK_MAP = fileinput.input()[0].strip()
109
+
110
+
disk = compact_disk(*read_disk(DISK_MAP))
111
+
print("Part 1:", disk_checksum(disk))
112
+
113
+
disk = compact_disk(*read_disk(DISK_MAP, part_2=True), part_2=True)
114
+
print("Part 2:", disk_checksum(disk))
115
+
116
+
117
+