advent-of-code/2024/day09/day9.py at master · buckley-w-david/advent-of-code · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
from collections import deque
from aoc_utils import *  # type: ignore

from aocd import get_data

data = get_data(year=2024, day=9, block=True)


def checksum_contiguous_disk(disk):
    checksum = 0
    for i, file_id in enumerate(disk):
        if file_id is None:
            return checksum
        checksum += i * file_id
    return checksum


def parse_one(data):
    blocks = [int(c) for c in data]
    disk = [None] * sum(blocks)
    state = True
    head = 0
    file_id = 0
    holes = deque([])
    files = {}
    for block in blocks:
        if state:
            disk[head : head + block] = [file_id] * block
            files[file_id] = list(range(head + block - 1, head - 1, -1))
            state = False
            file_id += 1
        else:
            holes.extend(range(head, head + block))
            state = True
        head += block
    return disk, holes, files


def part_one(data):
    disk, holes, files = parse_one(data)
    for file_id in list(reversed(files.keys())):
        positions = files[file_id]
        i = 0
        while holes and i < len(positions):
            if holes[0] > positions[i]:
                return checksum_contiguous_disk(disk)
            h = holes.popleft()
            disk[h] = file_id
            disk[positions[i]] = None
            i += 1
    return checksum_contiguous_disk(disk)


def parse_two(data):
    blocks = [int(c) for c in data]
    disk = [None] * sum(blocks)
    state = True
    head = 0
    file_id = 0
    holes = []
    files = {}
    for block in blocks:
        if state:
            disk[head : head + block] = [file_id] * block
            files[file_id] = (head, head + block)
            file_id += 1
        elif block != 0:
            holes.append((head, head + block))
        state = not state
        head += block
    return disk, holes, files


def checksum(disk):
    checksum = 0
    for i, file_id in enumerate(disk):
        if file_id is None:
            continue
        checksum += i * file_id
    return checksum


def part_two(data):
    disk, holes, files = parse_two(data)
    for file_id in list(reversed(files.keys())):
        start, end = files[file_id]
        length = end - start
        for i, (hs, he) in enumerate(holes):
            hl = he - hs
            if hs > start:
                break
            if length <= hl:
                disk[hs : hs + length] = [file_id] * length
                disk[start:end] = [None] * length
                if length == hl:
                    holes.pop(i)
                else:
                    holes[i] = (hs + length, he)
                before_start = None
                after_end = None
                # FIXME: O(n**2) is bad
                for i, (hs, he) in enumerate(reversed(holes)):
                    j = len(holes) - i - 1
                    if hs == end:
                        after_end = (j, (hs, he))
                    elif he == start:
                        before_start = (j, (hs, he))
                    if he < start:
                        break
                if before_start and after_end:
                    (i, (hs, _)) = before_start
                    (ri, (_, he)) = after_end
                    holes.pop(ri)
                elif before_start:
                    (i, (hs, _)) = before_start
                    he = end
                elif after_end:
                    (i, (_, he)) = after_end
                    hs = start
                else:
                    # This is technically incorrect, but the only time it matters is the very rightmost file
                    # Since we only ever try to move a file once, this doesn't actually come into play
                    break
                holes[i] = (hs, he)
                break

    return checksum(disk)


print(part_one(data))
print(part_two(data))