diff --git a/aoc2024/src/day11/python/solution.py b/aoc2024/src/day11/python/solution.py index 5192df2..23f93dd 100644 --- a/aoc2024/src/day11/python/solution.py +++ b/aoc2024/src/day11/python/solution.py @@ -11,57 +11,78 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +from collections import defaultdict -class Stone: - def __init__(self, value: int, next_stone=None): + +class _BinaryTree: + def __init__(self, value: int, parent=None, left=None, right=None): self.value: int = value - self.next: Stone | None = next_stone + self.parent: _BinaryTree | None = parent + self.left: _BinaryTree | None = left + self.right: _BinaryTree | None = right + + +def _has_even_num_digits(number: int) -> bool: + return len(str(number)) % 2 == 0 + + +def _parse(initial_state: str) -> tuple[_BinaryTree, ...]: + """Returns binary tree roots.""" + return tuple( + [_BinaryTree(value=int(stone)) for stone in initial_state.split()]) class Simulation: def __init__(self, initial_state: str): - self._first_stone = self._parse(initial_state) - - def _parse(self, initial_state: str) -> Stone: - stones = initial_state.split() - first = Stone(value=int(stones[0])) - prev = first - for i in range(1, len(stones)): - curr = Stone(value=int(stones[i])) - prev.next = curr - prev = curr - return first + self._roots = _parse(initial_state) + # Tracks value and descendants by iteration. + self._cache: dict[int, list[int, ...]] = defaultdict(list) - def blink(self): - curr = self._first_stone + def _update_cache(self, node, add_value): + # TODO: fix. + self._cache[node.value].append(add_value) + # Update ascendants + curr = node.parent while curr is not None: - if curr.value == 0: - curr.value = 1 - elif len(str(curr.value)) % 2 == 0: # Even num of digits: - half_digits = len(str(curr.value)) // 2 - # Break stone into 2: - first_value = str(curr.value)[:half_digits] - second_value = str(curr.value)[half_digits:] - second_stone = Stone(value=int(second_value), next_stone=curr.next) - curr.value = int(first_value) - curr.next = second_stone - curr = second_stone # Careful, should not handle the second stone. - else: - curr.value *= 2024 - curr = curr.next + values = self._cache[curr.value] + last_value = self._cache[curr.value][len(values) - 1] + self._cache[curr.value].append(last_value + add_value) + curr = curr.parent - def count_stones(self) -> int: - # Traverse stones from first. - num_stones = 0 - curr = self._first_stone - while curr is not None: - num_stones += 1 - curr = curr.next - return num_stones + def _iterate(self, node: _BinaryTree, num_iterations: int = 0) -> int: + """Returns the number of stones after iteration.""" + if num_iterations == 0: + return 1 # Leaf node + # TODO: use cache. + # if node.value in self._cache: + # return self._cache[node.value][num_iterations - 1] + if node.value == 0: + node.left = _BinaryTree(value=1, parent=node) + self._update_cache(node=node, add_value=1) + return self._iterate(node.left, num_iterations - 1) + elif _has_even_num_digits(node.value): + half_digits = len(str(node.value)) // 2 + # Break stone into 2: + left_node = _BinaryTree(value=int(str(node.value)[:half_digits]), + parent=node) + right_node = _BinaryTree(value=int(str(node.value)[half_digits:]), + parent=node) + node.left = left_node + node.right = right_node + self._update_cache(node=node, add_value=2) + return (self._iterate(node.left, num_iterations - 1) + + self._iterate(node.right, num_iterations - 1)) + else: + node.left = _BinaryTree(value=node.value * 2024, parent=node) + self._update_cache(node=node, add_value=1) + return self._iterate(node.left, num_iterations - 1) + + def simulate(self, num_iterations=0) -> int: + result = 0 + for root in self._roots: + result += self._iterate(root, num_iterations) + return result def count_stones(initial_state: str, blinks: int = 0) -> int: - simulation = Simulation(initial_state) - for _ in range(blinks): - simulation.blink() - return simulation.count_stones() + return Simulation(initial_state).simulate(num_iterations=blinks) diff --git a/aoc2024/test/day11/python/test_solution.py b/aoc2024/test/day11/python/test_solution.py index b1e9c80..32c448e 100644 --- a/aoc2024/test/day11/python/test_solution.py +++ b/aoc2024/test/day11/python/test_solution.py @@ -25,16 +25,25 @@ def __init__(self, *args, **kwargs): def test_part1_withOneBlink_counts(self): self.assertEqual(7, count_stones(initial_state='0 1 10 99 999', blinks=1)) + def test_part1_with3Blinks_counts(self): + self.assertEqual(5, count_stones(initial_state='125 17', blinks=3)) + def test_part1_withMoreBlinks_counts(self): - initial_state = '125 17' - self.assertEqual(22, count_stones(initial_state, blinks=6)) - self.assertEqual(55312, count_stones(initial_state, blinks=25)) + self.assertEqual(22, count_stones(initial_state='125 17', blinks=6)) + + def test_part1_withEvenMoreBlinks_counts(self): + self.assertEqual(55312, count_stones(initial_state='125 17', blinks=25)) def test_part1_withPuzzleInput_counts(self): self.assertEqual(203457, count_stones( initial_state='1 24596 0 740994 60 803 8918 9405859', blinks=25)) + # def test_part2_withPuzzleInput_counts(self): + # self.assertEqual(203457, count_stones( + # initial_state='1 24596 0 740994 60 803 8918 9405859', + # blinks=75)) + if __name__ == '__main__': unittest.main()