diff --git a/2024/src/aoc/days/day9.py b/2024/src/aoc/days/day9.py new file mode 100644 index 0000000..e65f1af --- /dev/null +++ b/2024/src/aoc/days/day9.py @@ -0,0 +1,107 @@ +import heapq + +from . import SeparateRunner + + +def file_checksum(file_id: int, start: int, length: int) -> int: + return file_id * length * (2 * start + length - 1) // 2 + + +class DayRunner(SeparateRunner): + @classmethod + def part1(cls, input: str) -> int: + files = [] + empty = [] + + pos = 0 + + for c in input.strip(): + val = int(c) + + if len(files) == len(empty): + files.append((pos, val)) + else: + empty.append((pos, val)) + + pos += val + + checksum = 0 + + for start, length in empty: + while files and length > 0: + file_start, file_len = files.pop() + if file_start < start: + files.append((file_start, file_len)) + break + + file_id = len(files) + + infill = min(file_len, length) + + checksum += file_checksum(file_id, start, infill) + start += infill + + if infill != file_len: + files.append((file_start, file_len - infill)) + + length -= infill + else: + continue + break + + for file_id, (file_start, file_len) in enumerate(files): + checksum += file_checksum(file_id, file_start, file_len) + + return checksum + + @classmethod + def part2(cls, input: str) -> int: + files = [] + empty = [[] for _ in range(10)] + + pos = 0 + + is_file = True + + for c in input.strip(): + val = int(c) + + if is_file: + files.append((pos, val)) + is_file = False + else: + # No need for heappush, as we're appending values in order + empty[val].append(pos) + is_file = True + + pos += val + + checksum = 0 + + while files: + start, length = files.pop() + file_id = len(files) + + best = None + best_heap = None + + for i, heap in enumerate(empty[length:]): + if not heap or heap[0] > start: + continue + + if best is None or best > heap[0]: + best = heap[0] + best_heap = i + length + + if best is None: + # No room to move left, count score at current position + checksum += file_checksum(file_id, start, length) + else: + checksum += file_checksum(file_id, best, length) + heapq.heappop(empty[best_heap]) + + if length < best_heap: + remainder = best_heap - length + heapq.heappush(empty[remainder], best + length) + + return checksum diff --git a/2024/tests/test_day09.py b/2024/tests/test_day09.py new file mode 100644 index 0000000..31c87c7 --- /dev/null +++ b/2024/tests/test_day09.py @@ -0,0 +1,11 @@ +from aoc.days.day9 import DayRunner + +SAMPLE = "2333133121414131402" + + +def test_sample_part1() -> None: + assert DayRunner.part1(SAMPLE) == 1928 + + +def test_sample_part2() -> None: + assert DayRunner.part2(SAMPLE) == 2858