Files
advent-of-code-2024/puzzles/9.py
2024-12-13 22:20:06 -08:00

302 lines
8.1 KiB
Python

import math
from itertools import count, pairwise
from typing import Iterable, NamedTuple, override
from unittest import TestCase
from more_itertools import chunked, ilen
from puzzles._solver import Solver
class LinkedList[T]:
__head: "LinkedListNode[T] | None"
__tail: "LinkedListNode[T] | None"
__slots__ = ("__head", "__tail")
@property
def head(self) -> "LinkedListNode[T] | None":
if not self.__head and self.__tail:
self.__head = self.__tail
while self.__head and self.__head.left:
self.__head = self.__head.left
return self.__head
@property
def tail(self) -> "LinkedListNode[T] | None":
if not self.__tail and self.__head:
self.__tail = self.__head
while self.__tail and self.__tail.right:
self.__tail = self.__tail.right
return self.__tail
def __init__(self, *values: Iterable[T]) -> None:
self.__head = None
self.__tail = None
for value in values:
self.append(value)
def append(self, value: T) -> None:
if tail := self.tail:
self.__tail = LinkedListNode(value, tail, None)
tail.right = self.__tail
else:
self.__tail = LinkedListNode(value, None, None)
def prepend(self, value: T) -> None:
if head := self.head:
self.__head = LinkedListNode(value, None, head)
head.left = self.__head
else:
self.__head = LinkedListNode(value, None, None)
def pop_tail(self) -> T | None:
if (tail := self.tail) is None:
raise IndexError("pop from empty LinkedList")
if tail.left:
tail.left.right = None
self.__tail = tail.left
tail.left = None
return tail.value
def pop_head(self) -> T:
if (head := self.head) is None:
raise IndexError("pop from empty LinkedList")
if head.right:
head.right.left = None
self.__head = head.right
head.right = None
return head.value
def copy(self) -> "LinkedList[T]":
return LinkedList(*self)
def __iter__(self) -> Iterable[T]:
node = self.head
while node is not None:
yield node.value
node = node.right
def __reversed__(self) -> Iterable[T]:
node = self.tail
while node is not None:
yield node.value
node = node.left
def __len__(self) -> int:
return ilen(self)
def __str__(self) -> str:
return f"LinkedList({', '.join(map(str, self))})"
class LinkedListNode[T]:
value: T
left: "LinkedListNode[T] | None"
right: "LinkedListNode[T] | None"
__slots__ = ("value", "left", "right")
def __init__(
self,
value: T,
left: "LinkedListNode[T] | None",
right: "LinkedListNode[T] | None",
) -> None:
self.value = value
self.left = left
self.right = right
def insert_left(self, value: T) -> None:
node = LinkedListNode(value, self.left, self)
if self.left:
self.left.right = node
self.left = node
def insert_right(self, value: T) -> None:
node = LinkedListNode(value, self, self.right)
if self.right:
self.right.left = node
self.right = node
def pop(self) -> T:
if self.left:
self.left.right = self.right
if self.right:
self.right.left = self.left
return self.value
class TestLinkedList(TestCase):
def test(self):
linked_list = LinkedList(1, 2, 3, 4, 5)
self.assertEqual(linked_list.head.value, 1)
self.assertEqual(linked_list.tail.value, 5)
self.assertListEqual(list(linked_list), [1, 2, 3, 4, 5])
self.assertListEqual(list(reversed(linked_list)), [5, 4, 3, 2, 1])
self.assertEqual(len(linked_list), 5)
self.assertEqual(linked_list.pop_tail(), 5)
self.assertEqual(linked_list.tail.value, 4)
self.assertEqual(linked_list.pop_head(), 1)
self.assertEqual(linked_list.head.value, 2)
linked_list.append(6)
self.assertEqual(linked_list.tail.value, 6)
linked_list.prepend(7)
self.assertEqual(linked_list.head.value, 7)
self.assertListEqual(list(linked_list), [7, 2, 3, 4, 6])
middle = linked_list.tail.left.left
self.assertEqual(middle.value, 3)
middle.insert_left(8)
self.assertListEqual(list(linked_list), [7, 2, 8, 3, 4, 6])
middle.insert_right(9)
self.assertListEqual(list(linked_list), [7, 2, 8, 3, 9, 4, 6])
self.assertEqual(middle.pop(), 3)
self.assertListEqual(list(linked_list), [7, 2, 8, 9, 4, 6])
class Block(NamedTuple):
id: int
position: int
size: int
def compact_drive(drive: LinkedList[Block]) -> LinkedList[Block]:
if not len(drive):
return drive
drive = drive.copy()
moving = drive.pop_tail()
cursor = drive.head
assert cursor is not None
while cursor.right is not None:
a = cursor.value
b = cursor.right.value
free_space = b.position - a.position - a.size
if not free_space:
cursor = cursor.right
continue
insertion = Block(
moving.id,
a.position + a.size,
min(free_space, moving.size),
)
cursor.insert_right(insertion)
cursor = cursor.right
if insertion.size < moving.size:
moving = Block(
moving.id,
moving.position,
moving.size - insertion.size,
)
else:
moving = drive.pop_tail()
cursor.insert_right(
Block(
moving.id,
cursor.value.position + cursor.value.size,
moving.size,
)
)
return drive
def compact_drive_without_fragmentation(drive: LinkedList[Block]) -> LinkedList[Block]:
if not len(drive):
return drive
drive = drive.copy()
moving = drive.tail
while moving:
cursor = drive.head
while cursor != moving:
a = cursor.value
b = cursor.right.value
free_space = b.position - a.position - a.size
if free_space >= moving.value.size:
moving.pop()
cursor.insert_right(
Block(
moving.value.id,
a.position + a.size,
moving.value.size,
)
)
break
cursor = cursor.right
moving = moving.left
return drive
def hash_drive(drive: LinkedList[Block]) -> int:
return sum(
block.id * block.size * (block.position * 2 + block.size - 1) // 2
for block in drive
)
def print_drive(drive: LinkedList[Block]) -> None:
snapshot = []
for a, b in pairwise(drive):
snapshot.append(str(a.id) * a.size)
snapshot.append("." * (b.position - a.position - a.size))
snapshot.append(str(drive.tail.value.id) * drive.tail.value.size)
print("".join(snapshot))
class DayNineSolver(Solver):
drive: LinkedList[Block]
@override
def __init__(self, puzzle_input: str):
self.drive = LinkedList()
snapshot = map(int, puzzle_input.strip())
cursor_position = 0
current_id = count()
for pair in chunked(snapshot, 2):
if data_size := pair[0]:
id = next(current_id)
self.drive.append(Block(id, cursor_position, data_size))
cursor_position += data_size
if len(pair) == 2:
cursor_position += pair[1]
@override
def solve_p1(self) -> int:
compacted_drive = compact_drive(self.drive)
return hash_drive(compacted_drive)
@override
def solve_p2(self) -> int:
compacted_drive = compact_drive_without_fragmentation(self.drive)
return hash_drive(compacted_drive)
class TestDayNineSolver(TestCase):
def test(self):
solver = DayNineSolver("2333133121414131402\n")
self.assertEqual(solver.solve_p1(), 1928)
self.assertEqual(solver.solve_p2(), 2858)