AoC2024.11 Change linked list to tree

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)

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()