AoC 2024 day 3 part 2

aoc2024/src/day03/python/solution.py

@@ -13,13 +13,46 @@
 # limitations under the License.
 import re
 from functools import reduce
+from typing import Sequence
+
+def _to_muls(matches: Sequence[str]) -> Sequence[tuple[int, ...]]:
+    # Ignore match[0] which contains the raw pair ('1','2') if mul(1,2).
+    # match[1] and match[2] capture the individual digits in pair (as string).
+    return tuple(tuple(map(int, (match[1], match[2]))) for match in matches)  # t:O(n)
+
+def _calculate(muls: Sequence[tuple[int, ...]]) -> int:
+    return sum(reduce(lambda a,b: a*b, mul) for mul in muls)  # t:O(n)
+
+def _find_muls(string: str) -> Sequence[str]:
+    return re.findall(r'mul\(((\d+),(\d+))\)', string)  # t:O(n)
 
 def calculate(input: str) -> int:
-    """Sums all the valid multiplications.
-    Example multiplication: mul(1,2)
     """
-    matches = re.findall(r'mul\(((\d+),(\d+))\)', input)
+    Sums all the valid multiplications.
+    Example valid multiplications: 'mul(1,2)', 'mul(123,456)'
+    """
+    matches = _find_muls(input)
     if matches is None:
         raise ValueError('No multiplications were found.')
-    muls = tuple(tuple(map(int, (match[1], match[2]))) for match in matches)  # O(n)
-    return sum(reduce(lambda a,b: a*b, mul) for mul in muls)  # O(n)
+    return _calculate(_to_muls(matches))
+
+def calculate2(input: str) -> int:
+    """
+    Sums all valid multiplications between do() and don't() operations.
+
+    At the beginning of the program, mul instructions are enabled.
+    Example: mul(1,1)do()mul(2,2)don't() -> 1*1+2*2=5
+    """
+    # Assumption: do() does not appear before don't(). Verified on my puzzle input.
+    result = 0
+    first_dont_pos = input.find('don\'t()')
+    result += calculate(input[:first_dont_pos])
+
+    sub_input = input[first_dont_pos+len('don\'t()'):]  # Trim beginning.
+    while sub_input.find('do()') != -1:
+        sub_input = sub_input[sub_input.find('do()'):] # Trim anything before first do().
+        dont_pos = sub_input.find('don\'t()')  # Find the next don't().
+        result += calculate(sub_input[:dont_pos])  # Calculate muls between do() and don't()
+        sub_input = sub_input[dont_pos+len('don\'t()'):]  # Trim anything before don't()
+    # Assumption: last operation is don't(). Verified on my puzzle input.
+    return result

aoc2024/test/day03/python/example2.txt

@@ -0,0 +1 @@
+xmul(2,4)&mul[3,7]!^don't()_mul(5,5)+mul(32,64](mul(11,8)undo()?mul(8,5))

aoc2024/test/day03/python/example3.txt

@@ -0,0 +1 @@
+xmul(2,4)&mul[3,7]!mul(1,1)^don't()_mul(5,5)+mul(32,64](mul(11,8)undo()?mul(8,5)andmul(1,1)don't())

aoc2024/test/day03/python/test_solution.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 import unittest
 
-from aoc2024.src.day03.python.solution import calculate
+from aoc2024.src.day03.python.solution import calculate, calculate2
 from common.python3.AdventOfCodeTestCase import AdventOfCodeTestCase
 
 
@@ -28,6 +28,15 @@ def test_part1_withExample_calculates(self):
   def test_part1_withPuzzleInput_calculates(self):
     self.assertEqual(170807108, calculate(self.input))
 
+  def test_part2_withExample_calculates(self):
+    self.assertEqual(48, calculate2(self.examples[1]))
+
+  def test_part2_withCraftedExample_calculates(self):
+    self.assertEqual(50, calculate2(self.examples[2]))
+
+  def test_part2_withPuzzleInput_calculates(self):
+    self.assertEqual(74838033, calculate2(self.input))
+
 
 if __name__ == '__main__':
   unittest.main()