JonathanGupton

Advent of Code 2024 - Day 09 - Python

Dec 9th, 2024 (edited)
126
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 3.52 KB | None | 0 0
  1. from dataclasses import dataclass
  2. from collections import deque
  3. from more_itertools import chunked
  4. from typing import Iterable
  5.  
  6.  
  7. @dataclass
  8. class Disk:
  9.     idx: int
  10.     files: int
  11.     free: int = 0
  12.  
  13.  
  14. def parse_data(fp: str) -> list[Disk]:
  15.     disk_space = []
  16.     with open(fp, "r") as f:
  17.         for idx, vals in enumerate(chunked(f.read().strip(), 2)):
  18.             disk_space.append(Disk(*map(int, (idx, *vals))))
  19.     return disk_space
  20.  
  21.  
  22. def checksum(files: Iterable[int | None]) -> int:
  23.     checksum_value = 0
  24.     for idx, val in enumerate(files):
  25.         checksum_value += idx * val
  26.     return checksum_value
  27.  
  28.  
  29. def defragment_disk_blocks(disk_space=list[Disk]) -> deque[int]:
  30.     output_queue = deque()
  31.     disk_space = deque(disk_space)
  32.     while disk_space:
  33.         left_file = disk_space.popleft()
  34.         for i in range(left_file.files):
  35.             output_queue.append(left_file.idx)
  36.         while disk_space and left_file.free:
  37.             right_file = disk_space.pop()
  38.             to_distribute = min(left_file.free, right_file.files)
  39.             for _ in range(to_distribute):
  40.                 output_queue.append(right_file.idx)
  41.             right_file.files -= to_distribute
  42.             left_file.free -= to_distribute
  43.             if right_file.files:
  44.                 disk_space.append(right_file)
  45.     return output_queue
  46.  
  47.  
  48. def defragment_disk_files(disk_space=list[Disk]):
  49.     disks = deque(disk_space)
  50.     hold = deque()
  51.     to_check = len(disk_space)
  52.     while to_check:
  53.         append = False
  54.         to_check -= 1
  55.         disk = disks.pop()
  56.         for i in range(len(disks)):
  57.             if disk.files <= disks[i].free:
  58.                 append = True
  59.                 break
  60.         if append:
  61.             disks[-1].free += disk.free + disk.files
  62.             disks[i].free, disk.free = 0, disks[i].free - disk.files
  63.             disks.insert(i + 1, disk)
  64.         else:
  65.             hold.appendleft(disk)
  66.     disks.extend(hold)
  67.     return disks
  68.  
  69.  
  70. def transform_disk_queue_to_blocks(disk_queue: deque[Disk]) -> list[int]:
  71.     blocks = []
  72.     for disk in disk_queue:
  73.         for _ in range(disk.files):
  74.             blocks.append(disk.idx)
  75.         for _ in range(disk.free):
  76.             blocks.append(0)
  77.     return blocks
  78.  
  79.  
  80. def example_a():
  81.     fp = "./example/day09-example01.txt"
  82.     data = parse_data(fp)
  83.     q = deque(map(int, "0099811188827773336446555566"))
  84.     defragmented = defragment_disk_blocks(data)
  85.     try:
  86.         assert q == defragmented
  87.     except AssertionError:
  88.         from itertools import zip_longest
  89.  
  90.         for idx, (l, r) in enumerate(zip_longest(q, defragmented)):
  91.             print(idx, l, r)
  92.  
  93.     checksum_value = checksum(defragmented)
  94.     assert checksum_value == 1928
  95.     print(checksum_value, "= 1928")
  96.  
  97.  
  98. def part_a():
  99.     fp = "./data/day09.txt"
  100.     disk_space = parse_data(fp)
  101.     defrag_disk = defragment_disk_blocks(disk_space)
  102.     cs = checksum(defrag_disk)
  103.     print(cs)
  104.  
  105.  
  106. def example_b():
  107.     fp = "./example/day09-example01.txt"
  108.     data = parse_data(fp)
  109.     defragmented = defragment_disk_files(data)
  110.     blocks = transform_disk_queue_to_blocks(defragmented)
  111.     cs = checksum(blocks)
  112.     print(cs, "= 2858")
  113.  
  114.  
  115. def part_b():
  116.     fp = "./data/day09.txt"
  117.     disk_space = parse_data(fp)
  118.     defragmented = defragment_disk_files(disk_space)
  119.     blocks = transform_disk_queue_to_blocks(defragmented)
  120.     cs = checksum(blocks)
  121.     print(cs)
  122.  
  123.  
  124. if __name__ == "__main__":
  125.     example_a()
  126.     part_a()
  127.     example_b()
  128.     part_b()
  129.  
Advertisement
Add Comment
Please, Sign In to add comment