302 lines
8.1 KiB
Python
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)
|