Restructure solver modules with classes and unittest suites
This commit is contained in:
355
puzzles/9.py
355
puzzles/9.py
@@ -1,26 +1,166 @@
|
||||
from itertools import islice, pairwise
|
||||
from typing import NamedTuple
|
||||
from itertools import count, pairwise
|
||||
import math
|
||||
from typing import Iterable, NamedTuple, override
|
||||
from unittest import TestCase
|
||||
|
||||
from more_itertools import chunked, sliced
|
||||
from more_itertools import chunked, ilen
|
||||
|
||||
test_input = """
|
||||
2333133121414131402
|
||||
"""
|
||||
|
||||
test_solution_p1 = 1928
|
||||
test_solution_p2 = 2858
|
||||
from puzzles._solver import Solver
|
||||
|
||||
|
||||
def solve_p1(puzzle_input: str) -> int:
|
||||
blocks = _parse_blocks(puzzle_input)
|
||||
compacted_blocks = _compact_blocks(blocks)
|
||||
return _hash_blocks(compacted_blocks)
|
||||
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))})"
|
||||
|
||||
|
||||
def solve_p2(puzzle_input: str) -> int:
|
||||
blocks = _parse_blocks(puzzle_input)
|
||||
compacted_blocks = _compact_blocks_no_fragmentation(blocks)
|
||||
return _hash_blocks(compacted_blocks)
|
||||
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):
|
||||
@@ -29,110 +169,133 @@ class Block(NamedTuple):
|
||||
size: int
|
||||
|
||||
|
||||
def _parse_blocks(puzzle_input: str) -> list[Block]:
|
||||
blocks: list[Block] = []
|
||||
cursor_position = 0
|
||||
current_id = 0
|
||||
def compact_drive(drive: LinkedList[Block]) -> LinkedList[Block]:
|
||||
if not len(drive):
|
||||
return drive
|
||||
|
||||
snapshot = map(int, puzzle_input.strip())
|
||||
for pair in chunked(snapshot, 2):
|
||||
if data_size := pair[0]:
|
||||
blocks.append(Block(current_id, cursor_position, data_size))
|
||||
cursor_position += data_size
|
||||
current_id += 1
|
||||
if len(pair) == 2:
|
||||
cursor_position += pair[1]
|
||||
drive = drive.copy()
|
||||
moving = drive.pop_tail()
|
||||
cursor = drive.head
|
||||
assert cursor is not None
|
||||
|
||||
return blocks
|
||||
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
|
||||
|
||||
def _compact_blocks(blocks: list[Block]) -> list[Block]:
|
||||
assert blocks
|
||||
blocks = blocks.copy()
|
||||
insertion = Block(
|
||||
moving.id,
|
||||
a.position + a.size,
|
||||
min(free_space, moving.size),
|
||||
)
|
||||
|
||||
compacted_blocks: list[Block] = []
|
||||
current_position = 0
|
||||
moving_block = blocks.pop()
|
||||
cursor.insert_right(insertion)
|
||||
cursor = cursor.right
|
||||
|
||||
while blocks:
|
||||
next_block = blocks.pop(0)
|
||||
free_space = next_block.position - current_position
|
||||
while free_space:
|
||||
assert moving_block.size > 0
|
||||
allocation = min(free_space, moving_block.size)
|
||||
compacted_blocks.append(
|
||||
Block(
|
||||
moving_block.id,
|
||||
current_position,
|
||||
allocation,
|
||||
)
|
||||
if insertion.size < moving.size:
|
||||
moving = Block(
|
||||
moving.id,
|
||||
moving.position,
|
||||
moving.size - insertion.size,
|
||||
)
|
||||
current_position += allocation
|
||||
free_space -= allocation
|
||||
if moving_block.size > allocation:
|
||||
moving_block = Block(
|
||||
moving_block.id,
|
||||
moving_block.position + allocation,
|
||||
moving_block.size - allocation,
|
||||
)
|
||||
elif blocks:
|
||||
moving_block = blocks.pop()
|
||||
else:
|
||||
break
|
||||
else:
|
||||
moving = drive.pop_tail()
|
||||
|
||||
compacted_blocks.append(next_block)
|
||||
current_position += next_block.size
|
||||
|
||||
compacted_blocks.append(
|
||||
cursor.insert_right(
|
||||
Block(
|
||||
moving_block.id,
|
||||
current_position,
|
||||
moving_block.size,
|
||||
moving.id,
|
||||
cursor.value.position + cursor.value.size,
|
||||
moving.size,
|
||||
)
|
||||
)
|
||||
|
||||
return compacted_blocks
|
||||
return drive
|
||||
|
||||
|
||||
def _compact_blocks_no_fragmentation(blocks: list[Block]) -> list[Block]:
|
||||
assert blocks
|
||||
blocks = blocks.copy()
|
||||
def compact_drive_without_fragmentation(drive: LinkedList[Block]) -> LinkedList[Block]:
|
||||
if not len(drive):
|
||||
return drive
|
||||
|
||||
compacted_blocks: list[Block] = []
|
||||
drive = drive.copy()
|
||||
moving = drive.tail
|
||||
|
||||
while len(blocks) > 1:
|
||||
moving_block = blocks[-1]
|
||||
for (_, a), (i, b) in pairwise(enumerate(blocks)):
|
||||
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_block.size:
|
||||
modified_block = Block(
|
||||
moving_block.id,
|
||||
a.position + a.size,
|
||||
moving_block.size,
|
||||
|
||||
if free_space >= moving.value.size:
|
||||
moving.pop()
|
||||
cursor.insert_right(
|
||||
Block(
|
||||
moving.value.id,
|
||||
a.position + a.size,
|
||||
moving.value.size,
|
||||
)
|
||||
)
|
||||
blocks.pop()
|
||||
blocks.insert(i, modified_block)
|
||||
break
|
||||
else:
|
||||
blocks.pop()
|
||||
compacted_blocks.append(moving_block)
|
||||
|
||||
compacted_blocks.append(blocks[0])
|
||||
cursor = cursor.right
|
||||
|
||||
return list(sorted(compacted_blocks, key=lambda block: block.position))
|
||||
moving = moving.left
|
||||
|
||||
return drive
|
||||
|
||||
|
||||
def _print_blocks(blocks: list[Block]) -> str:
|
||||
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(blocks):
|
||||
for a, b in pairwise(drive):
|
||||
snapshot.append(str(a.id) * a.size)
|
||||
snapshot.append("." * (b.position - a.position - a.size))
|
||||
snapshot.append(str(blocks[-1].id) * blocks[-1].size)
|
||||
snapshot.append(str(drive.tail.value.id) * drive.tail.value.size)
|
||||
print("".join(snapshot))
|
||||
|
||||
|
||||
def _hash_blocks(blocks: list[Block]) -> int:
|
||||
return sum(
|
||||
block.id * pos
|
||||
for block in blocks
|
||||
for pos in range(block.position, block.position + block.size)
|
||||
)
|
||||
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)
|
||||
|
Reference in New Issue
Block a user